{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Title: #Maximize Value of Function in a Ball Passing Game"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Difficulty: #Hard"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Category Title: #Algorithms"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Tag Slug: #bit-manipulation #array #dynamic-programming"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Name Translated: #位运算 #数组 #动态规划"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Solution Name: getMaxFunctionValue"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Translated Title: #在传球游戏中最大化函数值"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Translated Content:\n",
    "<p>给你一个长度为 <code>n</code>&nbsp;下标从 <strong>0</strong>&nbsp;开始的整数数组&nbsp;<code>receiver</code>&nbsp;和一个整数&nbsp;<code>k</code>&nbsp;。</p>\n",
    "\n",
    "<p>总共有&nbsp;<code>n</code>&nbsp;名玩家，玩家 <strong>编号</strong>&nbsp;互不相同，且为&nbsp;<code>[0, n - 1]</code>&nbsp;中的整数。这些玩家玩一个传球游戏，<code>receiver[i]</code>&nbsp;表示编号为 <code>i</code>&nbsp;的玩家会传球给编号为 <code>receiver[i]</code>&nbsp;的玩家。玩家可以传球给自己，也就是说&nbsp;<code>receiver[i]</code>&nbsp;可能等于&nbsp;<code>i</code>&nbsp;。</p>\n",
    "\n",
    "<p>你需要从 <code>n</code>&nbsp;名玩家中选择一名玩家作为游戏开始时唯一手中有球的玩家，球会被传 <strong>恰好</strong>&nbsp;<code>k</code>&nbsp;次。</p>\n",
    "\n",
    "<p>如果选择编号为 <code>x</code>&nbsp;的玩家作为开始玩家，定义函数&nbsp;<code>f(x)</code>&nbsp;表示从编号为&nbsp;<code>x</code>&nbsp;的玩家开始，<code>k</code>&nbsp;次传球内所有接触过球玩家的编号之&nbsp;<strong>和</strong>&nbsp;，如果有玩家多次触球，则 <strong>累加多次</strong>&nbsp;。换句话说，&nbsp;<code>f(x) = x + receiver[x] + receiver[receiver[x]] + ... + receiver<sup>(k)</sup>[x]</code>&nbsp;。</p>\n",
    "\n",
    "<p>你的任务时选择开始玩家 <code>x</code>&nbsp;，目的是<strong>&nbsp;最大化</strong>&nbsp;<code>f(x)</code>&nbsp;。</p>\n",
    "\n",
    "<p>请你返回函数的 <strong>最大值</strong>&nbsp;。</p>\n",
    "\n",
    "<p><strong>注意：</strong><code>receiver</code>&nbsp;可能含有重复元素。</p>\n",
    "\n",
    "<p>&nbsp;</p>\n",
    "\n",
    "<p><strong class=\"example\">示例 1：</strong></p>\n",
    "\n",
    "<table border=\"1\" cellspacing=\"3\" style=\"border-collapse: separate; text-align: center;\">\n",
    "\t<tbody>\n",
    "\t\t<tr>\n",
    "\t\t\t<th style=\"padding: 5px; border: 1px solid black;\">传递次数</th>\n",
    "\t\t\t<th style=\"padding: 5px; border: 1px solid black;\">传球者编号</th>\n",
    "\t\t\t<th style=\"padding: 5px; border: 1px solid black;\">接球者编号</th>\n",
    "\t\t\t<th style=\"padding: 5px; border: 1px solid black;\">x + 所有接球者编号</th>\n",
    "\t\t</tr>\n",
    "\t\t<tr>\n",
    "\t\t\t<td style=\"padding: 5px; border: 1px solid black;\">&nbsp;</td>\n",
    "\t\t\t<td style=\"padding: 5px; border: 1px solid black;\">&nbsp;</td>\n",
    "\t\t\t<td style=\"padding: 5px; border: 1px solid black;\">&nbsp;</td>\n",
    "\t\t\t<td style=\"padding: 5px; border: 1px solid black;\">2</td>\n",
    "\t\t</tr>\n",
    "\t\t<tr>\n",
    "\t\t\t<td style=\"padding: 5px; border: 1px solid black;\">1</td>\n",
    "\t\t\t<td style=\"padding: 5px; border: 1px solid black;\">2</td>\n",
    "\t\t\t<td style=\"padding: 5px; border: 1px solid black;\">1</td>\n",
    "\t\t\t<td style=\"padding: 5px; border: 1px solid black;\">3</td>\n",
    "\t\t</tr>\n",
    "\t\t<tr>\n",
    "\t\t\t<td style=\"padding: 5px; border: 1px solid black;\">2</td>\n",
    "\t\t\t<td style=\"padding: 5px; border: 1px solid black;\">1</td>\n",
    "\t\t\t<td style=\"padding: 5px; border: 1px solid black;\">0</td>\n",
    "\t\t\t<td style=\"padding: 5px; border: 1px solid black;\">3</td>\n",
    "\t\t</tr>\n",
    "\t\t<tr>\n",
    "\t\t\t<td style=\"padding: 5px; border: 1px solid black;\">3</td>\n",
    "\t\t\t<td style=\"padding: 5px; border: 1px solid black;\">0</td>\n",
    "\t\t\t<td style=\"padding: 5px; border: 1px solid black;\">2</td>\n",
    "\t\t\t<td style=\"padding: 5px; border: 1px solid black;\">5</td>\n",
    "\t\t</tr>\n",
    "\t\t<tr>\n",
    "\t\t\t<td style=\"padding: 5px; border: 1px solid black;\">4</td>\n",
    "\t\t\t<td style=\"padding: 5px; border: 1px solid black;\">2</td>\n",
    "\t\t\t<td style=\"padding: 5px; border: 1px solid black;\">1</td>\n",
    "\t\t\t<td style=\"padding: 5px; border: 1px solid black;\">6</td>\n",
    "\t\t</tr>\n",
    "\t</tbody>\n",
    "</table>\n",
    "\n",
    "<p>&nbsp;</p>\n",
    "\n",
    "<pre>\n",
    "<b>输入：</b>receiver = [2,0,1], k = 4\n",
    "<b>输出：</b>6\n",
    "<b>解释：</b>上表展示了从编号为 x = 2 开始的游戏过程。\n",
    "从表中可知，f(2) 等于 6 。\n",
    "6 是能得到最大的函数值。\n",
    "所以输出为 6 。\n",
    "</pre>\n",
    "\n",
    "<p><strong class=\"example\">示例 2：</strong></p>\n",
    "\n",
    "<table border=\"1\" cellspacing=\"3\" style=\"border-collapse: separate; text-align: center;\">\n",
    "\t<tbody>\n",
    "\t\t<tr>\n",
    "\t\t\t<th style=\"padding: 5px; border: 1px solid black;\">传递次数</th>\n",
    "\t\t\t<th style=\"padding: 5px; border: 1px solid black;\">传球者编号</th>\n",
    "\t\t\t<th style=\"padding: 5px; border: 1px solid black;\">接球者编号</th>\n",
    "\t\t\t<th style=\"padding: 5px; border: 1px solid black;\">x + 所有接球者编号</th>\n",
    "\t\t</tr>\n",
    "\t\t<tr>\n",
    "\t\t\t<td style=\"padding: 5px; border: 1px solid black;\">&nbsp;</td>\n",
    "\t\t\t<td style=\"padding: 5px; border: 1px solid black;\">&nbsp;</td>\n",
    "\t\t\t<td style=\"padding: 5px; border: 1px solid black;\">&nbsp;</td>\n",
    "\t\t\t<td style=\"padding: 5px; border: 1px solid black;\">4</td>\n",
    "\t\t</tr>\n",
    "\t\t<tr>\n",
    "\t\t\t<td style=\"padding: 5px; border: 1px solid black;\">1</td>\n",
    "\t\t\t<td style=\"padding: 5px; border: 1px solid black;\">4</td>\n",
    "\t\t\t<td style=\"padding: 5px; border: 1px solid black;\">3</td>\n",
    "\t\t\t<td style=\"padding: 5px; border: 1px solid black;\">7</td>\n",
    "\t\t</tr>\n",
    "\t\t<tr>\n",
    "\t\t\t<td style=\"padding: 5px; border: 1px solid black;\">2</td>\n",
    "\t\t\t<td style=\"padding: 5px; border: 1px solid black;\">3</td>\n",
    "\t\t\t<td style=\"padding: 5px; border: 1px solid black;\">2</td>\n",
    "\t\t\t<td style=\"padding: 5px; border: 1px solid black;\">9</td>\n",
    "\t\t</tr>\n",
    "\t\t<tr>\n",
    "\t\t\t<td style=\"padding: 5px; border: 1px solid black;\">3</td>\n",
    "\t\t\t<td style=\"padding: 5px; border: 1px solid black;\">2</td>\n",
    "\t\t\t<td style=\"padding: 5px; border: 1px solid black;\">1</td>\n",
    "\t\t\t<td style=\"padding: 5px; border: 1px solid black;\">10</td>\n",
    "\t\t</tr>\n",
    "\t</tbody>\n",
    "</table>\n",
    "\n",
    "<p>&nbsp;</p>\n",
    "\n",
    "<pre>\n",
    "<b>输入：</b>receiver = [1,1,1,2,3], k = 3\n",
    "<b>输出：</b>10\n",
    "<b>解释：</b>上表展示了从编号为 x = 4 开始的游戏过程。\n",
    "从表中可知，f(4) 等于 10 。\n",
    "10 是能得到最大的函数值。\n",
    "所以输出为 10 。\n",
    "</pre>\n",
    "\n",
    "<p>&nbsp;</p>\n",
    "\n",
    "<p><strong>提示：</strong></p>\n",
    "\n",
    "<ul>\n",
    "\t<li><code>1 &lt;= receiver.length == n &lt;= 10<sup>5</sup></code></li>\n",
    "\t<li><code>0 &lt;= receiver[i] &lt;= n - 1</code></li>\n",
    "\t<li><code>1 &lt;= k &lt;= 10<sup>10</sup></code></li>\n",
    "</ul>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Description: [maximize-value-of-function-in-a-ball-passing-game](https://leetcode.cn/problems/maximize-value-of-function-in-a-ball-passing-game/description/)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Solutions: [maximize-value-of-function-in-a-ball-passing-game](https://leetcode.cn/problems/maximize-value-of-function-in-a-ball-passing-game/solutions/)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_cases = ['[2,0,1]\\n4', '[1,1,1,2,3]\\n3']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def getMaxFunctionValue(self, receiver: List[int], k: int) -> int:\n",
    "        n = len(receiver)\n",
    "        inorder = [0] * n\n",
    "        for v in receiver:\n",
    "            inorder[v] += 1\n",
    "            \n",
    "        ans = 0\n",
    "        seen = [False] * n \n",
    "        c_info = [None] * n\n",
    "        for i in range(n):\n",
    "            if inorder[i] > 0:\n",
    "                continue\n",
    "                \n",
    "            cross, c_size, c_sum = -1, -1, -1\n",
    "            l, r = receiver[i], receiver[receiver[i]]\n",
    "            while l != r:\n",
    "                if c_info[l] is not None:\n",
    "                    break\n",
    "                l = receiver[l]\n",
    "                r = receiver[receiver[r]]\n",
    "            if c_info[l] is not None:\n",
    "                cross = l\n",
    "                c_size, c_sum = c_info[l]\n",
    "            else:\n",
    "                r = i \n",
    "                while l != r:\n",
    "                    l = receiver[l]\n",
    "                    r = receiver[r]\n",
    "                cross = l\n",
    "\n",
    "                j = receiver[cross]\n",
    "                c_size, c_sum = 1, cross\n",
    "                while j != cross:\n",
    "                    c_size += 1\n",
    "                    c_sum += j \n",
    "                    j = receiver[j]\n",
    "                c_info[cross] = (c_size, c_sum) \n",
    "            \n",
    "            cur_sum, cur_end = i, i \n",
    "            tmp_k = k\n",
    "            while tmp_k > 0 and cur_end != cross:\n",
    "                cur_end = receiver[cur_end]\n",
    "                cur_sum += cur_end \n",
    "                tmp_k -= 1\n",
    "            cur_sum += c_sum * (tmp_k // c_size)\n",
    "            tmp_k %= c_size \n",
    "            while tmp_k > 0:\n",
    "                cur_end = receiver[cur_end]\n",
    "                cur_sum += cur_end \n",
    "                tmp_k -= 1\n",
    "                \n",
    "            cur_start = i\n",
    "            while not seen[cur_start]:\n",
    "                seen[cur_start] = True \n",
    "                ans = max(ans, cur_sum)\n",
    "                cur_sum = cur_sum - cur_start + receiver[cur_end] \n",
    "                cur_start, cur_end = receiver[cur_start], receiver[cur_end] \n",
    "        \n",
    "        for i in range(n):\n",
    "            if seen[i]:\n",
    "                continue\n",
    "            j = receiver[i]\n",
    "            c_size, c_sum = 1, i \n",
    "            while j != i:\n",
    "                c_size += 1\n",
    "                c_sum += j \n",
    "                j = receiver[j]\n",
    "            tmp_k = k \n",
    "            cur_sum = i + c_sum * (tmp_k // c_size)\n",
    "            tmp_k %= c_size \n",
    "            cur_end = i \n",
    "            while tmp_k > 0:\n",
    "                cur_end = receiver[cur_end]\n",
    "                cur_sum += cur_end \n",
    "                tmp_k -= 1\n",
    "            cur_start = i\n",
    "            while not seen[cur_start]:\n",
    "                seen[cur_start] = True \n",
    "                ans = max(ans, cur_sum)\n",
    "                cur_sum = cur_sum - cur_start + receiver[cur_end] \n",
    "                cur_start, cur_end = receiver[cur_start], receiver[cur_end] \n",
    "            \n",
    "        return ans "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def getMaxFunctionValue(self, receiver: List[int], k: int) -> int:\n",
    "        n = len(receiver)\n",
    "        link = []\n",
    "        index = [-1] * n\n",
    "        step = [0] * n\n",
    "        nextPoint = [-1] * n\n",
    "        deg = [0] * n\n",
    "        for i in range(n):\n",
    "            deg[receiver[i]] += 1\n",
    "\n",
    "        for d in (0,1):\n",
    "            for i in range(n):\n",
    "                if deg[i]!=d:continue\n",
    "                start = len(link)\n",
    "                while index[i] == -1:\n",
    "                    deg[i]-=1\n",
    "                    index[i] = len(link)\n",
    "                    link.append(i)\n",
    "                    i = receiver[i]\n",
    "                stones = [start, index[i], len(link)] if index[i] >= start else [start, len(link)]\n",
    "                for s, e in pairwise(stones):\n",
    "                    for j in range(s, e):\n",
    "                        step[j] = e - j\n",
    "                        nextPoint[j] = index[i]\n",
    "\n",
    "        presum=list(accumulate(link))\n",
    "        def f(i, k):\n",
    "            cur = 0\n",
    "            while k >= step[i]:\n",
    "                cur += presum[i + step[i] - 1] - presum[i] + link[i]\n",
    "                k -= step[i]\n",
    "                i = nextPoint[i]\n",
    "                if i == nextPoint[i]:\n",
    "                    T = presum[i + step[i] - 1] - presum[i] + link[i]\n",
    "                    cur += (k // step[i]) * T\n",
    "                    k %= step[i]\n",
    "            # k<step[i]\n",
    "            cur += presum[i + k] - presum[i] + link[i]\n",
    "            return cur,link[i+k]\n",
    "\n",
    "        ans = 0\n",
    "        i = 0\n",
    "        while i < n:\n",
    "            cur, knext = f(i, k)\n",
    "            ans = max(ans, cur)\n",
    "            while i + 1 < n and link[i + 1] == receiver[link[i]]:\n",
    "                knext = receiver[knext]\n",
    "                cur += knext - link[i]\n",
    "                ans = max(ans, cur)\n",
    "                i += 1\n",
    "            i += 1\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def getMaxFunctionValue(self, r: List[int], k: int) -> int:\n",
    "        n = len(r)\n",
    "        p = [i for i in range(n)]\n",
    "        q = p.copy()\n",
    "        u = r.copy()\n",
    "        v = r.copy()\n",
    "        \n",
    "        def mul(a, b, c, d):\n",
    "            p = [c[a[i]] for i in range(n)]\n",
    "            q = [b[i] + d[a[i]] for i in range(n)]\n",
    "            return p, q\n",
    "        \n",
    "        for i in range(k.bit_length()):\n",
    "            if k & (1 << i) > 0:\n",
    "                p, q = mul(p, q, u, v)\n",
    "            u, v = mul(u, v, u, v)\n",
    "        \n",
    "        return max(q)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def getMaxFunctionValue(self, receiver: List[int], k: int) -> int:\n",
    "        n = len(receiver)\n",
    "        # 赋值：初始位置以及起点的得分\n",
    "        ans = list(range(n))\n",
    "        pos = list(range(n))\n",
    "        # 赋值：走 2 ^ 0 步时的结束位置和得分\n",
    "        ans_2_power = receiver\n",
    "        pos_2_power = receiver\n",
    "        for i in range(34):\n",
    "            if k >> i & 1:\n",
    "                # 继续往前走 2 ^ i 步\n",
    "                ans = [ans[i] + ans_2_power[pos[i]] for i in range(n)]\n",
    "                pos = [pos[pos_2_power[i]] for i in range(n)]\n",
    "            # 从 2 ^ i 步变成 2 ^ (i + 1) 步\n",
    "            ans_2_power = [ans_2_power[i] + ans_2_power[pos_2_power[i]] for i in range(n)]\n",
    "            pos_2_power = [pos_2_power[x] for x in pos_2_power]\n",
    "        return max(ans)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# 2836.在传球游戏中最大化函数值\n",
    "# https://leetcode.cn/problems/maximize-value-of-function-in-a-ball-passing-game/description/\n",
    "# lang=python\n",
    "\n",
    "# 解法:图论，带环链表\n",
    "from collections import defaultdict, deque\n",
    "from typing import List\n",
    "\n",
    "\n",
    "class Solution:\n",
    "    def getMaxFunctionValue(self, receiver: List[int], k: int) -> int:\n",
    "        k += 1\n",
    "        n = len(receiver)\n",
    "        \n",
    "        # 第一步找出图里的所有环，顺便计算从环上开始的最大值（用滑动窗口）\n",
    "        ans = 0\n",
    "        prefixs = defaultdict(list)\n",
    "        vis = [0] * n\n",
    "        loc = [0] * n\n",
    "        for x in range(n):\n",
    "            if vis[x]: continue\n",
    "            seen = {}\n",
    "            ls = list()\n",
    "            while 1:\n",
    "                vis[x] = 1\n",
    "                seen[x] = len(ls)\n",
    "                ls.append(x)\n",
    "                x = receiver[x]\n",
    "                if vis[x]:\n",
    "                    if x not in seen: break\n",
    "                    ls = ls[seen[x]:]\n",
    "                    m = len(ls)\n",
    "                    prefix = [0]\n",
    "                    for idx, j in enumerate(ls):\n",
    "                        prefixs[j] = prefix\n",
    "                        loc[j] = idx\n",
    "                    ls = ls + ls\n",
    "                    for j in range(2 * m):\n",
    "                        prefix.append(prefix[-1] + ls[j])\n",
    "                    tmp1 = prefix[m] * (k // m)\n",
    "                    r = k % m\n",
    "                    if r:\n",
    "                        tmp2 = cnt = sum(ls[:r])\n",
    "                        for i in range(r, m + r):\n",
    "                            cnt += ls[i] - ls[i - r]\n",
    "                            tmp2 = max(tmp2, cnt)\n",
    "                        tmp1 += tmp2\n",
    "                    ans = max(ans, tmp1)\n",
    "                    break\n",
    "        \n",
    "        # 第二部，从入度为0的点开始走，类似滑动窗口遍历每个点出发能得到结果，然后当然是取最大值~\n",
    "        indegree = [0] * n\n",
    "        for v in receiver:\n",
    "            indegree[v] += 1\n",
    "        \n",
    "        for i in range(n):\n",
    "            if indegree[i]: continue\n",
    "            ptr = i\n",
    "            q = deque()\n",
    "            cnt = 0\n",
    "            while ptr not in prefixs:\n",
    "                q.append(ptr)\n",
    "                cnt += ptr\n",
    "                if len(q) == k:\n",
    "                    ans = max(ans, cnt)\n",
    "                    cnt -= q.popleft()\n",
    "                ptr = receiver[ptr]\n",
    "            prefix, idx = prefixs[ptr], loc[ptr]\n",
    "            m = len(prefix) // 2\n",
    "            kk = k - len(q)\n",
    "            while q:\n",
    "                ans = max(ans, cnt + prefix[idx + kk % m] - prefix[idx] + (kk // m) * prefix[m])\n",
    "                cnt -= q.popleft()\n",
    "                kk += 1\n",
    "        \n",
    "        return ans\n",
    "\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    import time\n",
    "    \n",
    "    t0 = time.time()\n",
    "    \n",
    "    s = Solution()\n",
    "    func_name = dir(s)[-1]\n",
    "    func = getattr(s, func_name)\n",
    "    \n",
    "    print(func(receiver=[2, 0, 1], k=4))\n",
    "    # 输出：6\n",
    "    print(func(receiver=[1, 1, 1, 2, 3], k=3))\n",
    "    # 输出：10\n",
    "    print(func([3, 4, 1, 2, 2], 10))\n",
    "    # 输出：27\n",
    "    print(func([2, 1, 1], 4))\n",
    "    # 输出：6\n",
    "    print(func([2, 3, 3, 1], 15))\n",
    "    # 输出：33\n",
    "    print(func([1, 0, 1], 2))\n",
    "    # 输出：3\n",
    "    print(func([3, 2, 2, 1, 0], 2))\n",
    "    # 输出：7\n",
    "    \n",
    "    t1 = time.time()\n",
    "    print(t1 - t0)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def getMaxFunctionValue(self, r: List[int], k: int) -> int:\n",
    "        n = len(r)\n",
    "        lst = defaultdict(list)\n",
    "        deg = [0 for _ in range(n)]\n",
    "        for i in range(n):\n",
    "            fa,nxt = i,r[i]\n",
    "            deg[nxt] += 1\n",
    "            lst[nxt].append(fa)\n",
    "        vis = set()\n",
    "        vv = []\n",
    "        cy = {}\n",
    "        # 在环上绕圈\n",
    "        def maxincy(v):\n",
    "            nonlocal k\n",
    "            n = len(v)\n",
    "            cc = k // n\n",
    "            left = k % n\n",
    "            pre = [0]\n",
    "            suf = [0]\n",
    "            for c in v:\n",
    "                pre.append(pre[-1] + c)\n",
    "            for c in v[::-1]:\n",
    "                suf.append(suf[-1] + c)\n",
    "            suf = suf[::-1]\n",
    "            ans = 0\n",
    "            for i in range(n):\n",
    "                # 从i开始\n",
    "                tmp = cc * pre[-1] + v[i]\n",
    "                right = n - 1 - i\n",
    "                \n",
    "                if right >= left:\n",
    "                \n",
    "                   tmp += pre[i + left + 1] - pre[i + 1]\n",
    "                else:\n",
    "                    tt = left - right\n",
    "                    tmp += suf[i + 1] + pre[tt]\n",
    "                ans = max(ans,tmp)\n",
    "            return ans\n",
    "                    \n",
    "        def getv(v1,v2):\n",
    "            n1,n2 = len(v1),len(v2)\n",
    "            nonlocal k\n",
    "            ans = 0\n",
    "            p1 = [0]\n",
    "            for c in v1:\n",
    "                p1.append(p1[-1] + c)\n",
    "            p2 = [0]\n",
    "            for c in v2:\n",
    "                p2.append(p2[-1] + c)\n",
    "            for i in range(n1):\n",
    "                if k > n1 - 1 - i:\n",
    "                    left = k - (n1 - 1 - i)\n",
    "                    tmp = p1[-1] - p1[i]\n",
    "                else:\n",
    "                    tmp = p1[i + k + 1] - p1[i]\n",
    "                    ans = max(ans,tmp)\n",
    "                    continue\n",
    "                left -= 1\n",
    "                tmp += v2[0]\n",
    "                cyint,lt = left // n2,left % n2\n",
    "                tmp += cyint * p2[-1] + p2[lt + 1] - p2[1]\n",
    "                ans = max(ans,tmp)\n",
    "            return ans\n",
    "            \n",
    "        ans = 0\n",
    "        for i in range(n):\n",
    "            if deg[i] == 0 and i not in vis:\n",
    "                tmp = [i]\n",
    "                v = {i}\n",
    "                nxt = r[i]\n",
    "                while nxt not in v and nxt not in cy:\n",
    "                    tmp.append(nxt)\n",
    "                    v.add(nxt)\n",
    "                    nxt = r[nxt]\n",
    "                if nxt in cy:\n",
    "                   cv = cy[nxt]\n",
    "                   ans = max(ans,maxincy(cv))\n",
    "                   ans = max(ans,getv(tmp,cv))\n",
    "                else:\n",
    "                    id = tmp.index(nxt)\n",
    "                    cy[nxt] = tmp[id:]\n",
    "                    ans = max(ans,maxincy(tmp[id:]))\n",
    "                    ans = max(ans,getv(tmp[:id],tmp[id:]))                \n",
    "                for j in v:\n",
    "                    vis.add(j)\n",
    "    \n",
    "        for i in range(n):\n",
    "            if i not in vis:\n",
    "                tmp = [i]\n",
    "                v = {i}\n",
    "                nxt = r[i]\n",
    "                while nxt not in v:\n",
    "                    v.add(nxt)\n",
    "                    tmp.append(nxt)\n",
    "                    nxt = r[nxt]\n",
    "                for j in v:\n",
    "                    vis.add(j)\n",
    "                ans = max(ans,maxincy(tmp))\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def getMaxFunctionValue(self, receiver: List[int], k: int) -> int:\n",
    "        k += 1\n",
    "        n = len(receiver)\n",
    "        before = [[] for i in range(n)]\n",
    "        for i in range(len(receiver)):\n",
    "            before[receiver[i]].append(i)\n",
    "        test = [len(before[i]) for i in range(n)]\n",
    "        que = []\n",
    "        vis = set()\n",
    "        for i in range(n):\n",
    "            if test[i] == 0:\n",
    "                que.append(i)\n",
    "                vis.add(i)\n",
    "        while que:\n",
    "            x = que.pop()\n",
    "            test[receiver[x]] -= 1\n",
    "            if test[receiver[x]] == 0:\n",
    "                que.append(receiver[x])\n",
    "                vis.add(receiver[x])\n",
    "        ref = {}\n",
    "        for i in range(n):\n",
    "            if i not in vis:\n",
    "                p = [i]\n",
    "                x = i\n",
    "                res = i\n",
    "                while receiver[x] != i:\n",
    "                    x = receiver[x]\n",
    "                    p.append(x)\n",
    "                    res += x\n",
    "                for x in p:\n",
    "                    ref[x] = [len(p), res]\n",
    "                    vis.add(x)\n",
    "        ans = 0\n",
    "        tovis = []\n",
    "        for i in range(n):\n",
    "            if len(before[i]) == 0:\n",
    "                tovis.append(i)\n",
    "        vis = set()\n",
    "        for i in tovis:\n",
    "            x = i\n",
    "            num = 0\n",
    "            res = 0\n",
    "            que = deque()\n",
    "            while num < k:\n",
    "                if x in ref:\n",
    "                    t = (k - num) // ref[x][0]\n",
    "                    res += t * ref[x][1]\n",
    "                    num += t * ref[x][0]\n",
    "                if num == k:\n",
    "                    break\n",
    "                que.append(x)\n",
    "                res += x\n",
    "                num += 1\n",
    "                x = receiver[x]\n",
    "            ans = max(ans, res)\n",
    "            if not que:\n",
    "                vis.add(i)\n",
    "            else:\n",
    "                while que:\n",
    "                    x = que[-1]\n",
    "                    r = que.popleft()\n",
    "                    if r in vis:\n",
    "                        break\n",
    "                    vis.add(r)\n",
    "                    res -= r\n",
    "                    res += receiver[x]\n",
    "                    ans = max(ans, res)\n",
    "                    que.append(receiver[x])\n",
    "        for i in range(n):\n",
    "            if i not in vis:\n",
    "                x = i\n",
    "                res = 0\n",
    "                num = 0\n",
    "                que = deque()\n",
    "                if x in ref:\n",
    "                    t = k // ref[x][0]\n",
    "                    res += t * ref[x][1]\n",
    "                    num += t * ref[x][0]\n",
    "                if num == k:\n",
    "                    ans = max(res, ans)\n",
    "                    vis.add(i)\n",
    "                    continue\n",
    "                while num < k:\n",
    "                    num += 1\n",
    "                    que.append(x)\n",
    "                    res += x\n",
    "                    x = receiver[x]\n",
    "                while que:\n",
    "                    x = que[-1]\n",
    "                    r = que.popleft()\n",
    "                    if r in vis:\n",
    "                        break\n",
    "                    vis.add(r)\n",
    "                    res -= r\n",
    "                    res += receiver[x]\n",
    "                    ans = max(ans, res)\n",
    "                    que.append(receiver[x])\n",
    "        return ans\n",
    "                    \n",
    "                \n",
    "                    \n",
    "                    \n",
    "                    \n",
    "            \n",
    "                \n",
    "                \n",
    "                    \n",
    "                \n",
    "                \n",
    "                \n",
    "            \n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def getMaxFunctionValue(self, receiver: List[int], k: int) -> int:\n",
    "\n",
    "        f = set(receiver)\n",
    "        n = len(receiver)\n",
    "        c = [-1]*n\n",
    "        p = defaultdict(list)\n",
    "        sp = defaultdict(list)\n",
    "        kk = 0\n",
    "        # find loop\n",
    "        def dfs(i):  \n",
    "            nonlocal kk\n",
    "            visit = set()\n",
    "            while c[i] < 0 and i not in visit:\n",
    "                visit.add(i)\n",
    "                i = receiver[i]\n",
    "            \n",
    "            if c[i] < 0:\n",
    "                c[i] = kk\n",
    "                p[kk].append(i)\n",
    "                m = receiver[i]\n",
    "                while m != i:\n",
    "                    p[kk].append(m)\n",
    "                    m = receiver[m]\n",
    "                kk += 1\n",
    "                                  \n",
    "            for j in visit:\n",
    "                c[j] = c[i]\n",
    "                if j not in f: sp[c[j]].append(j)\n",
    "                    \n",
    "        for i in range(n):\n",
    "            if c[i] >= 0: continue\n",
    "            dfs(i)\n",
    "        \n",
    "        score = [-1] * n\n",
    "        # get score\n",
    "        def gets(i):\n",
    "            w = {}\n",
    "            s = [0] * (len(p[i])+1)\n",
    "            for idx, j in enumerate(p[i]):\n",
    "                w[j] = idx\n",
    "                s[idx+1] = s[idx]+j\n",
    "            \n",
    "            if len(sp[i]) > 0:\n",
    "                for j in sp[i]:\n",
    "                    score[j] = j\n",
    "                    end = receiver[j]\n",
    "                    m = 0\n",
    "                    while m < k:\n",
    "                        if end in w: break\n",
    "                        score[j] += end\n",
    "                        end = receiver[end]\n",
    "                        m += 1\n",
    "                    if m < k:\n",
    "                        res = k - m\n",
    "                        score[j] += (res // len(p[i])) * s[-1]\n",
    "                        res %= len(p[i])\n",
    "                        l, r = w[end], w[end] + res\n",
    "                        if r <= len(p[i]):\n",
    "                            score[j] += s[r] - s[l]\n",
    "                        else:\n",
    "                            r %= len(p[i])\n",
    "                            score[j] += s[-1] - s[l] + s[r]\n",
    "                        end = p[i][r % len(p[i])]\n",
    "                        \n",
    "                    now = receiver[j]\n",
    "                    while score[now] == -1:\n",
    "                        score[now] = score[j] - j + end\n",
    "                        end = receiver[end]\n",
    "                        j = now\n",
    "                        now = receiver[now]\n",
    "            else:\n",
    "                #print(p[i])\n",
    "                j = p[i][0]\n",
    "                score[j] = ((k+1) // len(p[i])) * s[-1]\n",
    "                res = (k+1) % len(p[i])\n",
    "                score[j] += s[res]\n",
    "                end = p[i][res]\n",
    "                now = receiver[j]\n",
    "                while now and score[now] == -1:\n",
    "                    score[now] = score[j] - j + end\n",
    "                    end = receiver[end]\n",
    "                    j = now\n",
    "                    now = receiver[now]\n",
    "                #print(score)\n",
    "                \n",
    "        for i in range(kk):\n",
    "            gets(i)\n",
    "        \n",
    "        return max(score)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def getMaxFunctionValue(self, A: List[int], k: int) -> int:\n",
    "        n = len(A)\n",
    "        indeg = [0] * n\n",
    "        for x in A: indeg[x] += 1\n",
    "        q = [x for x in range(n) if indeg[x] == 0]\n",
    "        ss = q.copy()\n",
    "        while q:\n",
    "            nq = []\n",
    "            for u in q:\n",
    "                v = A[u]\n",
    "                # rg[v].append(u)\n",
    "                indeg[v] -= 1\n",
    "                if indeg[v] == 0: nq.append(v)\n",
    "            q = nq\n",
    "\n",
    "        seen = [False] * n\n",
    "        C = []\n",
    "        book = [None] * n\n",
    "        for u in range(n):\n",
    "            if indeg[u] == 0 or seen[u]: continue\n",
    "            \n",
    "            S = []\n",
    "            while not seen[u]:\n",
    "                seen[u] = True\n",
    "                S.append(u)\n",
    "                u = A[u]\n",
    "            C.append((S, sum(S), len(S), list(accumulate(S+S, initial=0))))\n",
    "            p = len(C) - 1\n",
    "            for i, u in enumerate(S):\n",
    "                book[u] = (p, i)\n",
    "        # print(indeg)\n",
    "        # print(C)\n",
    "        # print(book)\n",
    "        \n",
    "        res = 0\n",
    "        for S, tot, sz, P in C:\n",
    "            ks = k // sz * tot\n",
    "            rk = k % sz\n",
    "            for i, u in enumerate(S):\n",
    "                res = max(res, ks + u + P[i+rk+1] - P[i+1])\n",
    "                # print(ks + u + P[i+rk+1] - P[i+1])\n",
    "        \n",
    "        seen = [False] * n\n",
    "        \n",
    "        def gao(u: int):\n",
    "            B = []\n",
    "            while not book[u]:\n",
    "                B.append(u)\n",
    "                u = A[u]\n",
    "            p, i = book[u]\n",
    "            S, tot, sz, _ = C[p]\n",
    "            S = S[i:] + S[:i]\n",
    "            P = list(accumulate(S+S, initial=0))\n",
    "            PB = list(accumulate(B, initial=0))\n",
    "            nonlocal res\n",
    "            for i, u in enumerate(B):\n",
    "                if i + k < len(B):\n",
    "                    res = max(res, u + PB[i+k+1] - PB[i+1])\n",
    "                else:\n",
    "                    tk = k - (len(B) - i)\n",
    "                    ks = tk // sz * tot + PB[-1] - PB[i+1]\n",
    "                    rk = tk % sz\n",
    "                    res = max(res, u + ks + P[rk + 1] - P[0])\n",
    "            \n",
    "        \n",
    "        for u in ss:\n",
    "            gao(u)\n",
    "                \n",
    "        return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def getMaxFunctionValue(self, receiver: List[int], k: int) -> int:\n",
    "        n = len(receiver)\n",
    "        k += 1\n",
    "        indeg = [0] * n\n",
    "        for i, t in enumerate(receiver):\n",
    "            indeg[t] += 1\n",
    "        \n",
    "        q = deque(i for i, v in enumerate(indeg) if v == 0)\n",
    "        g = [[] for _ in range(n)]\n",
    "        \n",
    "        while q:\n",
    "            x = q.popleft()\n",
    "            y = receiver[x]\n",
    "            g[y].append(x)\n",
    "            indeg[y] -= 1\n",
    "            if indeg[y] == 0:\n",
    "                q.append(y)\n",
    "\n",
    "        seen = [False] * n\n",
    "        ans = 0\n",
    "\n",
    "        def dfs(pres: list, tail: list, x: int):\n",
    "            nonlocal ans\n",
    "            tail.append(x)\n",
    "            if len(tail) >= k:\n",
    "                ans = max(ans, sum(tail[-k:]))\n",
    "            else:\n",
    "                res = sum(tail)\n",
    "                d, m = divmod(k - len(tail), len(pres))\n",
    "                res += d * pres[-1]\n",
    "                if m: res += pres[m - 1]\n",
    "                ans = max(ans, res)\n",
    "            for y in g[x]:\n",
    "                dfs(pres, tail, y)\n",
    "            tail.pop()\n",
    "\n",
    "        \n",
    "        for i in range(n):\n",
    "            if indeg[i] == 0 or seen[i]: continue\n",
    "            path = []\n",
    "            s = 0\n",
    "            while not seen[i]:\n",
    "                seen[i] = True\n",
    "                path.append(i)\n",
    "                s += i\n",
    "                i = receiver[i]\n",
    "            l = len(path)\n",
    "            d, m = divmod(k, l)\n",
    "            res = d * s\n",
    "            if m:\n",
    "                mx = curr = sum(path[:m])\n",
    "                for j in range(1, l):\n",
    "                    curr += path[(j + m - 1) % l] - path[j - 1]\n",
    "                    mx = max(mx, curr)\n",
    "                res += mx\n",
    "            ans = max(ans, res)\n",
    "\n",
    "            for i, x in enumerate(path):\n",
    "                if not g[x]: continue\n",
    "                suff = path[i:] + path[:i]\n",
    "                # suff.reverse()\n",
    "                pres = list(accumulate(suff))\n",
    "                for y in g[x]:\n",
    "                    dfs(pres, [], y)\n",
    "\n",
    "        return ans\n",
    "\n",
    "\n",
    "\n",
    "# class Solution:\n",
    "#     def getMaxFunctionValue(self, receiver: List[int], k: int) -> int:\n",
    "#         n = len(receiver)\n",
    "#         ans = list(range(n))\n",
    "#         cur = list(range(n))\n",
    "#         f = receiver.copy()\n",
    "#         g = receiver.copy()\n",
    "#         i = 0\n",
    "#         while k:\n",
    "#             if (k >> i) & 1:\n",
    "#                 for j in range(n):\n",
    "#                     ans[j] += g[cur[j]]\n",
    "#                 c_cur = [0] * n\n",
    "#                 for j in range(n):\n",
    "#                     c_cur[j] = f[cur[j]]\n",
    "#                 cur = c_cur\n",
    "#                 k -= 1 << i\n",
    "#             ff = [0] * n\n",
    "#             gg = [0] * n\n",
    "#             for j in range(n):\n",
    "#                 ff[j] = f[f[j]]\n",
    "#                 gg[j] = g[j] + g[f[j]]\n",
    "#             f = ff\n",
    "#             g = gg\n",
    "#             i += 1\n",
    "#         return max(ans)\n",
    "\n",
    "# class Solution:\n",
    "#     def getMaxFunctionValue(self, r: List[int], k: int) -> int:\n",
    "#         ans = 0\n",
    "#         a = [None] * len(r)\n",
    "#         h = {}\n",
    "#         for i in range(len(r)):\n",
    "#             if a[i]:\n",
    "#                 continue\n",
    "#             b = [[0, 0] for _ in r]\n",
    "#             c = {i:-1}\n",
    "#             j = i\n",
    "#             b[i][1] = i\n",
    "#             while r[j] not in c and a[r[j]] is None:\n",
    "#                 c[r[j]] = j\n",
    "#                 b[r[j]][0] = b[j][0] + 1\n",
    "#                 b[r[j]][1] = b[j][1] + r[j]\n",
    "#                 j = r[j]\n",
    "#             if r[j] in c:\n",
    "#                 flag = c[r[j]]\n",
    "#                 cost = b[j][1] - b[r[j]][1] + r[j]\n",
    "#                 le = b[j][0] - b[r[j]][0] + 1\n",
    "#                 rukou = r[j]\n",
    "#                 h[rukou] = [0] * (le + 1)\n",
    "#                 h[rukou][-1] = cost\n",
    "#                 kk = 0\n",
    "#                 a[rukou] = [0, 0, cost, rukou, le]\n",
    "#                 while j != flag:\n",
    "#                     a[j] = list(a[r[j]])\n",
    "#                     a[j][0] += 1\n",
    "#                     a[j][1] += j\n",
    "#                     h[rukou][le - kk - 1] = h[rukou][le - kk] - j\n",
    "#                     kk += 1\n",
    "#                     j = c[j]\n",
    "#                 a[rukou] = [0, 0, cost, rukou, le]\n",
    "#             while j != -1:\n",
    "#                 a[j] = list(a[r[j]])\n",
    "#                 a[j][0] += 1\n",
    "#                 a[j][1] += j\n",
    "#                 j = c[j]\n",
    "#         print(a)\n",
    "#         print(h)\n",
    "#         for i in range(len(r)):\n",
    "#             # print(i, k)\n",
    "#             t = 0\n",
    "#             if a[i][0] > k:\n",
    "#                 print(a)\n",
    "#                 t += i\n",
    "#                 j = i\n",
    "#                 for _ in range(k):\n",
    "#                     j = r[j]\n",
    "#                     t += j\n",
    "#                     # if i == 4: print(t)\n",
    "#             else:\n",
    "#                 t += a[i][1] + ((k - a[i][0]) // a[i][4]) * a[i][2]\n",
    "#                 j = a[i][3]\n",
    "#                 t += h[j][(k - a[i][0]) % a[i][4] + 1]\n",
    "#             if t > ans:\n",
    "#                 ans = t\n",
    "#         return ans\n",
    "            \n",
    "            \n",
    "                \n",
    "                \n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def getMaxFunctionValue(self, receiver: List[int], k: int) -> int:\n",
    "        n=len(receiver)\n",
    "        used=[0]*n\n",
    "        groups=[[]]\n",
    "        gid=uid=0\n",
    "        chd=[[] for _ in range(n)]\n",
    "        for i in range(n):\n",
    "            if used[i]!=0:\n",
    "                continue\n",
    "            uid-=1\n",
    "            nd=i\n",
    "            while 1:\n",
    "                if used[nd]==0:\n",
    "                    used[nd]=uid\n",
    "                    chd[receiver[nd]].append(nd)\n",
    "                    nd=receiver[nd]\n",
    "                elif used[nd]!=uid:\n",
    "                    break\n",
    "                else:\n",
    "                    groups.append([])\n",
    "                    nd1=nd\n",
    "                    gid+=1\n",
    "                    while 1:\n",
    "                        groups[-1].append(nd1)\n",
    "                        del chd[nd1][-1]\n",
    "                        used[nd1]=gid\n",
    "                        nd1=receiver[nd1]\n",
    "                        if nd1==nd:\n",
    "                            break\n",
    "                    break\n",
    "        \n",
    "        def cal(gid):\n",
    "            rt=0\n",
    "            gt=groups[gid]\n",
    "            n=len(gt)\n",
    "            g=[0]*n\n",
    "            for i,a in enumerate(gt):\n",
    "                g[i]=a+(0 if i==0 else g[i-1])\n",
    "            def f(p,ln):\n",
    "                if ln<=0:\n",
    "                    return 0\n",
    "                a,b=divmod(ln,n)\n",
    "                rt=g[-1]*a\n",
    "                if b==0:\n",
    "                    pass\n",
    "                elif p+b<=n:\n",
    "                    rt+=g[p+b-1]-(0 if p==0 else g[p-1])\n",
    "                else:\n",
    "                    rt+=g[-1]-(0 if p==0 else g[p-1])+g[p-1+b-n]\n",
    "                return rt\n",
    "            \n",
    "            def f1(root,nd,lt,s,pb):\n",
    "                lt.append(nd)\n",
    "                s+=nd\n",
    "                ln=len(lt)\n",
    "                if ln-pb>k+1:\n",
    "                    s-=lt[pb]\n",
    "                    pb+=1\n",
    "                rt=s+f(root,k+1-ln)\n",
    "                for ci in chd[nd]:\n",
    "                    rt=max(rt,f1(root,ci,lt,s,pb))\n",
    "                del lt[-1]\n",
    "                return rt\n",
    "\n",
    "            for i,nd in enumerate(gt):\n",
    "                rt=max(rt,f(i,k+1))\n",
    "                # print(i,nd,f(i,k+1))\n",
    "                for ci in chd[nd]:\n",
    "                    rt=max(rt,f1(i,ci,[],0,0))\n",
    "                    # print(i,nd,ci,f1(i,ci,[],0,0))\n",
    "            return rt\n",
    "        \n",
    "        res=0\n",
    "        for i in range(1,gid+1):\n",
    "            res=max(res,cal(i))\n",
    "        return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def getMaxFunctionValue(self, receiver: List[int], k: int) -> int:\n",
    "        k += 1\n",
    "        n = len(receiver)\n",
    "        indeg = [0] * n\n",
    "        adj = [[] for _ in range(n)]\n",
    "        for u, v in enumerate(receiver):\n",
    "            adj[v].append(u)\n",
    "            indeg[v] += 1\n",
    "        \n",
    "        q = deque(u for u in range(n) if indeg[u] == 0)\n",
    "        while q:\n",
    "            u = q.popleft()\n",
    "            v = receiver[u]\n",
    "            indeg[v] -= 1\n",
    "            if indeg[v] == 0:\n",
    "                q.append(v)\n",
    "\n",
    "        f = [0] * n\n",
    "        g = [0] * n\n",
    "\n",
    "        # 环上的点\n",
    "        rings = collections.defaultdict(lambda:[])\n",
    "        # 某个环节上节点开始遍历的路径\n",
    "        paths = collections.defaultdict(lambda:[])\n",
    "\n",
    "        # 从环上节点反向遍历\n",
    "        def rdfs(u: int, dep: int, lasti: int, s: int, r: int):\n",
    "            ring = rings[r]\n",
    "            sz = len(ring)\n",
    "            if dep > 0:\n",
    "                g[u] = s\n",
    "            for v in adj[u]:\n",
    "                if indeg[v] != 0: continue\n",
    "                if dep < k:\n",
    "                    paths[r].append(v)\n",
    "                    rdfs(v, dep + 1, (lasti - 1 + sz) % sz, s + v - ring[lasti], r)\n",
    "                    paths[r].pop()\n",
    "                else:\n",
    "                    paths[r].append(v)\n",
    "                    rdfs(v, dep + 1, (lasti - 1 + sz) % sz, s + v - paths[r][-k-1], r)\n",
    "                    paths[r].pop()\n",
    "        \n",
    "        for u in range(n):\n",
    "            if indeg[u] == 1:\n",
    "                ring = []\n",
    "                x = u\n",
    "                while True:\n",
    "                    indeg[x] = -1\n",
    "                    ring.append(x)\n",
    "                    x = receiver[x]\n",
    "                    if x == u:\n",
    "                        break\n",
    "                m = len(ring)\n",
    "                prefix = list(accumulate(ring, initial=0))\n",
    "                for i, u in enumerate(ring):\n",
    "                    s = prefix[m] * (k // m)\n",
    "                    j = (i + k % m) % m\n",
    "                    if j > i:\n",
    "                        s += (prefix[j] - prefix[i])\n",
    "                    elif j < i:\n",
    "                        s += (prefix[m] - prefix[i] + prefix[j])\n",
    "                    \n",
    "                    f[u] = s\n",
    "\n",
    "                for i, u in enumerate(ring):\n",
    "                    if len(adj[u]) > 1:\n",
    "                        rings[u] = ring\n",
    "                        paths[u] = [u]\n",
    "                        lasti = (i - 1 + k % m) % m\n",
    "                        rdfs(u, 0, lasti, f[u], u)\n",
    "        return max(g[i] + f[i] for i in range(n))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def getMaxFunctionValue(self, receiver: List[int], k: int) -> int:\n",
    "        deg = [0] * len(receiver)\n",
    "        par = [[] for _ in range(len(receiver))]\n",
    "        for i, n in enumerate(receiver):\n",
    "            deg[n] += 1\n",
    "            par[n].append(i)\n",
    "        \n",
    "        q = []\n",
    "        for i, n in enumerate(deg):\n",
    "            if n == 0:\n",
    "                q.append(i)\n",
    "        \n",
    "        while q:\n",
    "            i = q.pop()\n",
    "            deg[receiver[i]] -= 1\n",
    "            if deg[receiver[i]] == 0:\n",
    "                q.append(receiver[i])\n",
    "        \n",
    "        cycle = [-1] * len(receiver)\n",
    "        idx = [-1] * len(receiver)\n",
    "        prefix = []\n",
    "        j = 0\n",
    "        for i, n in enumerate(deg):\n",
    "            if n == 0 or cycle[i] != -1:\n",
    "                continue\n",
    "            prefix.append([i])\n",
    "            idx[i] = 0\n",
    "            cycle[i] = j\n",
    "            nxt = receiver[i]\n",
    "            while nxt != i:\n",
    "                cycle[nxt] = j\n",
    "                prefix[-1].append(prefix[-1][-1] + nxt)\n",
    "                idx[nxt] = len(prefix[-1]) -1\n",
    "                nxt = receiver[nxt]\n",
    "            j += 1\n",
    "        \n",
    "        acc = [0]\n",
    "        def dfs(root, node):\n",
    "            acc.append(acc[-1] + node)\n",
    "            cycle_out = acc[-1] - (acc[0] if k >= len(acc)-2 else acc[len(acc)-k-2])\n",
    "            cycle_in = max(0, (k-len(acc)+2) // len(prefix[cycle[root]])) * prefix[cycle[root]][-1]\n",
    "            remain_num = max(0, k-len(acc)+2) % len(prefix[cycle[root]])\n",
    "            remain = (prefix[cycle[root]][idx[root] + remain_num] if idx[root] + remain_num < len(prefix[cycle[root]]) else prefix[cycle[root]][-1] + prefix[cycle[root]][(idx[root] + remain_num) % len(prefix[cycle[root]])] ) - prefix[cycle[root]][idx[root]]\n",
    "            ret = cycle_in + cycle_out + remain\n",
    "            for nxt in par[node]:\n",
    "                if deg[nxt] == 0:\n",
    "                    ret = max(ret, dfs(root, nxt))\n",
    "            acc.pop()\n",
    "            return ret\n",
    "        \n",
    "        ret = 0\n",
    "        for i, n in enumerate(deg):\n",
    "            \n",
    "            if n != 0:\n",
    "                ret = max(ret, dfs(i, i))\n",
    "        return ret"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def getMaxFunctionValue(self, receiver: List[int], k: int) -> int:\n",
    "        n = len(receiver)\n",
    "\n",
    "        # 找环\n",
    "        cycles: List[List[int]] = []\n",
    "        visited = [False] * n\n",
    "        in_cycles = [False] * n\n",
    "        for i in range(n):\n",
    "            path = []\n",
    "            u = i\n",
    "            while not visited[u]:\n",
    "                visited[u] = True\n",
    "                path.append(u)\n",
    "                u = receiver[u]\n",
    "            if u in path:\n",
    "                cycles.append(path[path.index(u):])\n",
    "                for root in cycles[-1]:\n",
    "                    in_cycles[root] = True\n",
    "\n",
    "        # 反向图\n",
    "        g = [[] for _ in range(n)]\n",
    "        for u in range(n):\n",
    "            g[receiver[u]].append(u)\n",
    "\n",
    "        # 通过 dfs 确定每个点到环后还需走多少次, 以及特判 k 次后不能到环的点的答案\n",
    "        tree: List[List[Tuple[int, int]]] = [[] for _ in range(n)]\n",
    "        ans = [0] * n\n",
    "        for cycle in cycles:\n",
    "            for root in cycle:\n",
    "                stack = []\n",
    "                # u: 结点, s: 到根点和\n",
    "                def dfs(u: int, s: int):\n",
    "                    s += u\n",
    "                    if len(stack) < k:\n",
    "                        tree[root].append((u, k - len(stack)))\n",
    "                        ans[u] += s\n",
    "                    else:\n",
    "                        ans[u] = s - stack[-k]\n",
    "                    stack.append(s - u)\n",
    "                    for v in g[u]:\n",
    "                        if not in_cycles[v]:\n",
    "                            dfs(v, s)\n",
    "                    stack.pop()\n",
    "                dfs(root, 0)\n",
    "        \n",
    "        # 枚举每个环, 使用前缀和统计环上的贡献\n",
    "        for cycle in cycles:\n",
    "            prefix = cycle.copy()\n",
    "            for i in range(1, len(prefix)):\n",
    "                prefix[i] += prefix[i - 1]\n",
    "            for i in range(len(cycle)):\n",
    "                for u, d in tree[cycle[i]]:\n",
    "                    ans[u] += d // len(cycle) * prefix[-1]\n",
    "                    d %= len(cycle)\n",
    "                    if i + d < len(cycle):\n",
    "                        ans[u] += prefix[i + d] - prefix[i]\n",
    "                    else:\n",
    "                        ans[u] += prefix[-1] - prefix[i] + prefix[i + d - len(cycle)]\n",
    "        \n",
    "        return max(ans)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "from collections import defaultdict\n",
    "\n",
    "class UnionFind:\n",
    "    def __init__(self, n):\n",
    "        self.ids = [i for i in range(n)]\n",
    "        self.h = [1]*n\n",
    "        self.count = n\n",
    "    \n",
    "    def find(self, p):\n",
    "        if self.ids[p] == p:\n",
    "            return p\n",
    "        return self.find(self.ids[p])\n",
    "    \n",
    "    def union(self, p, q):\n",
    "        p = self.find(p)\n",
    "        q = self.find(q)\n",
    "        if p != q:\n",
    "            if self.h[p] < self.h[q]:\n",
    "                p, q = q, p\n",
    "            self.ids[q] = p\n",
    "            if self.h[p] == self.h[q]:\n",
    "                self.h[p] += 1\n",
    "            self.count -= 1\n",
    "\n",
    "class Solution:\n",
    "    ans = 0\n",
    "    def getMaxFunctionValue(self, receiver: List[int], k: int) -> int:\n",
    "        uf = UnionFind(len(receiver))\n",
    "        for i in range(len(receiver)):\n",
    "            uf.union(receiver[i], i)\n",
    "        ns, r2s = defaultdict(list), defaultdict(list)\n",
    "        for i in range(len(receiver)):\n",
    "            ns[uf.find(i)].append(i)\n",
    "            r2s[receiver[i]].append(i)\n",
    "        \n",
    "        vis = set()\n",
    "        def dfs(_k, ci, c, cp, i, _ss):\n",
    "            if len(_ss) > _k:\n",
    "                _s = _ss[-1]-(_ss[len(_ss)-_k-2] if len(_ss) > _k+1 else 0)\n",
    "                self.ans = max(self.ans, _s)\n",
    "            else:\n",
    "                d = _k-len(_ss)\n",
    "                t, cj = d//len(c), (ci+d)%len(c)\n",
    "                if ci <= cj:\n",
    "                    _s = cp[cj]-cp[ci]+c[ci]\n",
    "                else:\n",
    "                    _s = cp[-1]+cp[cj]-cp[ci]+c[ci]\n",
    "                _s += t*cp[-1]+_ss[-1]\n",
    "                self.ans = max(self.ans, _s)\n",
    "            vis.add(i)\n",
    "            for j in r2s[i]:\n",
    "                if j in vis: continue\n",
    "                _ss.append(_ss[-1]+j)\n",
    "                dfs(_k, ci, c, cp, j, _ss)\n",
    "                _ss.pop()\n",
    "\n",
    "        for vs in ns.values():\n",
    "            i, _vis = vs[0], set()\n",
    "            while i not in _vis:\n",
    "                _vis.add(i)\n",
    "                i = receiver[i]\n",
    "            cur, ps = [], []\n",
    "            i2c = dict()\n",
    "            while i not in vis:\n",
    "                cur.append(i)\n",
    "                i2c[i] = len(cur)-1  \n",
    "                if len(ps) == 0:\n",
    "                    ps.append(i)\n",
    "                else:\n",
    "                    ps.append(ps[-1]+i)\n",
    "                vis.add(i)\n",
    "                i = receiver[i]\n",
    "            for i in range(len(cur)):\n",
    "                t, j = k//len(cur), (i+k)%len(cur)\n",
    "                if i <= j:\n",
    "                    _s = ps[j]-ps[i]+cur[i]\n",
    "                else:\n",
    "                    _s = ps[-1]+ps[j]-ps[i]+cur[i]\n",
    "                _s += t*ps[-1]\n",
    "                self.ans = max(self.ans, _s)\n",
    "                for j in r2s[cur[i]]:\n",
    "                    if j in vis: continue\n",
    "                    dfs(k, i2c[receiver[j]], cur, ps, j, [j])\n",
    "        \n",
    "        return self.ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def getMaxFunctionValue(self, receiver: List[int], k: int) -> int:\n",
    "        n = len(receiver)\n",
    "        # 建图\n",
    "        g = [[] for _ in range(n)]\n",
    "        rg = [[] for _ in range(n)]\n",
    "        indegree = [0] * n\n",
    "        for i in range(n):\n",
    "            g[i].append(receiver[i])\n",
    "            indegree[receiver[i]] += 1\n",
    "            rg[receiver[i]].append(i)\n",
    "\n",
    "        # 在环上的点和不在环上的点\n",
    "        q = deque([i for i, v in enumerate(indegree) if v == 0])\n",
    "        while q:\n",
    "            node = q.popleft()\n",
    "            for ch in g[node]:\n",
    "                indegree[ch] -= 1\n",
    "                if indegree[ch] == 0:\n",
    "                    q.append(ch)\n",
    "\n",
    "        # 找不在环上的点到第一个在环上的点的和以及距离\n",
    "        dis = [[] for _ in range(n)]\n",
    "        ans = [0] * n\n",
    "        for i in range(n):\n",
    "            if indegree[i] == 0:\n",
    "                continue\n",
    "            st = []\n",
    "\n",
    "            def dfs(node, cnt):\n",
    "                cnt += node\n",
    "                if len(st) < k:\n",
    "                    dis[node] = (i, k - len(st))\n",
    "                    ans[node] += cnt\n",
    "                else:\n",
    "                    dis[node] = (i, 0)\n",
    "                    ans[node] = cnt - st[-k]\n",
    "                \n",
    "                st.append(cnt - node)\n",
    "                for ch in rg[node]:\n",
    "                    if indegree[ch] == 0:\n",
    "                        dfs(ch, cnt)\n",
    "                st.pop()\n",
    "            \n",
    "            dfs(i, 0)\n",
    "\n",
    "        dic = defaultdict(list)\n",
    "        d = {}\n",
    "        pos = {}\n",
    "        for i in range(n):\n",
    "            if indegree[i] == 0:\n",
    "                continue\n",
    "            node = i\n",
    "            res = []\n",
    "            while indegree[node]:\n",
    "                d[node] = i\n",
    "                pos[node] = len(res)\n",
    "                res.append(node)\n",
    "                indegree[node] = 0\n",
    "                node = receiver[node]\n",
    "            dic[i] = res\n",
    "\n",
    "        ps = defaultdict(list)\n",
    "        for x, v in dic.items():\n",
    "            ps[x] = list(accumulate(v + v, initial=0))\n",
    "        \n",
    "        def check(node, q):\n",
    "            i = d[node]\n",
    "            nums = dic[i]\n",
    "            m = len(nums)\n",
    "            idx = (pos[node] + 1) % m\n",
    "            total = ps[i][m]\n",
    "            t, re = divmod(q, m)\n",
    "            return t * total + ps[i][idx + re] - ps[i][idx]\n",
    "        \n",
    "        for i in range(n):\n",
    "            ans[i] += check(*dis[i])\n",
    "\n",
    "        return max(ans)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def getMaxFunctionValue(self, receiver: List[int], k: int) -> int:\n",
    "        n = len(receiver)\n",
    "        pre_sum = receiver\n",
    "        pre_id = receiver\n",
    "        valid_pow2s=[pre_sum] if k&1>0 else []\n",
    "        ids_pow2s=[pre_id] if k&1>0 else []\n",
    "        step=2\n",
    "        while step<=k:\n",
    "            cur_ids = [pre_id[pre_id[i]] for i in range(n)]\n",
    "            cur_step_sum = [pre_sum[i]+pre_sum[pre_id[i]] for i in range(n)]\n",
    "            if k&step>0:\n",
    "                valid_pow2s.append(cur_step_sum)\n",
    "                ids_pow2s.append(cur_ids)\n",
    "            pre_id = cur_ids\n",
    "            pre_sum = cur_step_sum\n",
    "            step<<=1\n",
    "        # print(valid_pow2s)\n",
    "        # print(ids_pow2s)\n",
    "        res = list(range(n))\n",
    "        cur_id = list(range(n))\n",
    "        for ids,sums in zip(reversed(ids_pow2s), reversed(valid_pow2s)):\n",
    "            res = [res[i] + sums[cur_id[i]] for i in range(n)]\n",
    "            cur_id = [ids[cur_id[i]] for i in range(n)]\n",
    "        return max(res)\n",
    "            \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "\n",
    "class Solution:\n",
    "    def getMaxFunctionValue(self, receiver: List[int], k: int) -> int:\n",
    "        k+=1\n",
    "        n=len(receiver)\n",
    "        g=[[] for _ in range(n)]\n",
    "        for i in range(n):\n",
    "            g[receiver[i]].append(i)\n",
    "\n",
    "        cyc=[]\n",
    "        incyc=[False]*n\n",
    "        vis=[False]*n\n",
    "        comp=[False]*n\n",
    "        def find(i):\n",
    "            if comp[i]:return\n",
    "            if vis[i]:\n",
    "                cc=[i]\n",
    "                j=i\n",
    "                incyc[i]=True\n",
    "                while receiver[j]!=i:\n",
    "                    cc.append(receiver[j])\n",
    "                    j=receiver[j]\n",
    "                    incyc[j]=True\n",
    "                cyc.append(cc)\n",
    "                vis[i] = True\n",
    "                comp[i] = True\n",
    "                return\n",
    "            vis[i] = True\n",
    "            find(receiver[i])\n",
    "            comp[i] = True\n",
    "\n",
    "        for i in range(n):\n",
    "            find(i)\n",
    "\n",
    "        res=0\n",
    "        for cc in cyc:\n",
    "            cc=cc[::-1]\n",
    "            r=k//len(cc)\n",
    "            s=sum(cc)*r\n",
    "            j=(k-1)%len(cc)\n",
    "            if k%len(cc)!=0:\n",
    "                s+=sum(cc[:j+1])\n",
    "            for i in range(len(cc)):\n",
    "                res=max(res,s)\n",
    "\n",
    "                stack=[]\n",
    "                stack_sum=[0]\n",
    "                def dfs(i,cyc_sum,end):\n",
    "                    nonlocal res\n",
    "                    stack.append(end)\n",
    "                    stack_sum.append(stack_sum[-1]+end)\n",
    "                    if len(stack)>=k:\n",
    "                        res=max(res,stack_sum[len(stack)]-stack_sum[len(stack)-k])\n",
    "                    else:\n",
    "                        res=max(res,stack_sum[-1]+cyc_sum-cc[i])\n",
    "                    for ne in g[end]:\n",
    "                        dfs((i+1)%len(cc),cyc_sum-cc[i],ne)\n",
    "                    stack.pop()\n",
    "                    stack_sum.pop()\n",
    "                for ne in g[cc[j]]:\n",
    "                    if incyc[ne]: continue\n",
    "                    dfs(i,s,ne)\n",
    "\n",
    "                s-=cc[i]\n",
    "                j=(j+1)%len(cc)\n",
    "                s+=cc[j]\n",
    "        return res\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def getMaxFunctionValue(self, receiver: List[int], k: int) -> int:\n",
    "        k = k + 1\n",
    "        n = len(receiver)\n",
    "        de = [0] * n\n",
    "        for i in range(n):\n",
    "            de[receiver[i]] += 1\n",
    "        def dfs(p, dic, arr):\n",
    "            if p in dic:\n",
    "                return dic[p]\n",
    "            dic[p] = len(arr)\n",
    "            vi.add(p)\n",
    "            arr.append(p)\n",
    "            return dfs(receiver[p], dic, arr)\n",
    "        # print(de)\n",
    "        vi = set()\n",
    "        def go(arr, m):\n",
    "            mn = len(arr)\n",
    "            ans = 0\n",
    "            m2 = mn - m\n",
    "            arr2 = arr[m:]\n",
    "            if k <= mn:\n",
    "                s = sum(arr[:k])\n",
    "                e = k\n",
    "                ans = s\n",
    "                for i in range(mn):\n",
    "                    if e == mn:\n",
    "                        e = m\n",
    "                    s -= arr[i]\n",
    "                    # print(e, arr, m2, mn, m)\n",
    "                    s += arr[e]\n",
    "                    e += 1\n",
    "                    ans = max(ans, s)\n",
    "                return ans\n",
    "            else:\n",
    "                left = k - m\n",
    "                s1 = sum(arr[:m])\n",
    "                s2 = sum(arr2)\n",
    "                s = s1\n",
    "                s += s2 * (left // m2)\n",
    "                e = left % m2\n",
    "                s += sum(arr2[:e])\n",
    "                ans = s\n",
    "                # print(m2, s1, s2, s, e)\n",
    "                for i in range(mn):\n",
    "                    s -= arr[i]\n",
    "                    s += arr2[e]\n",
    "                    e += 1\n",
    "                    e %= m2\n",
    "                    ans = max(ans, s)\n",
    "                return ans\n",
    "        res = 0\n",
    "        for i in range(n):\n",
    "            if de[i] == 0:\n",
    "                arr = []\n",
    "                dic = {}\n",
    "                idx = dfs(i, dic, arr)\n",
    "                # print(arr, idx)\n",
    "                res = max(res, go(arr, idx))\n",
    "        for i in range(n):\n",
    "            if i not in vi:\n",
    "                arr = []\n",
    "                dic = {}\n",
    "                idx = dfs(i, dic, arr)\n",
    "                # print(arr, idx)\n",
    "                res = max(res, go(arr, idx))\n",
    "        return res\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "def mm(r, k):\n",
    "        g = defaultdict(list)\n",
    "        d = defaultdict(int)\n",
    "        fa = defaultdict(set)\n",
    "        for j in range(len(r)):\n",
    "            g[j].append(r[j]) \n",
    "            d[r[j]] += 1 \n",
    "            fa[r[j]].add(j)\n",
    "        s = [j for j in range(len(r)) if d[j]==0] \n",
    "        v = set()\n",
    "        while s:\n",
    "            tmp = []\n",
    "            for j in s:\n",
    "                v.add(j)\n",
    "                for kk in g[j]:\n",
    "                    d[kk] -= 1\n",
    "                    if d[kk]==0:\n",
    "                        tmp.append(kk) \n",
    "            s = tmp  \n",
    "        n = len(r)\n",
    "        c = [j for j in range(n) if j not in v] \n",
    "        v = {}\n",
    "       \n",
    "        def f(x):\n",
    "            st = x\n",
    "            tmp = 0\n",
    "            ix = 0 \n",
    "            dic = {}\n",
    "            pre = {}\n",
    "            while d[x]:\n",
    "                st = g[st][0]\n",
    "                d[st] -= 1\n",
    "                tmp += st \n",
    "                dic[st] = ix \n",
    "                pre[ix] = tmp \n",
    "                ix+=1\n",
    "#                 print(d[x])\n",
    "            mod = ix \n",
    "            # print(pre)\n",
    "            def gg(u, m): #u是真实的u\n",
    "                uj = u \n",
    "                u = dic[u]\n",
    "                nu = (u + m)%mod \n",
    "                cir = (u + m)//mod \n",
    "                ans = pre[nu] - pre[u] + pre[ix-1]*cir + uj \n",
    "                return ans   \n",
    "            # # print(gg(1, 3))\n",
    "            # print(k)\n",
    "            for j in dic:\n",
    "                v[j] = gg(j, k) \n",
    "            # print(v, gg(1, 3))\n",
    "            stack = []\n",
    "            def dfs(x, use):\n",
    "                use += x \n",
    "                stack.append(use)  \n",
    "                if x not in v:\n",
    "                    if len(stack)<=(k+1):\n",
    "                        jump = len(stack)-1 \n",
    "                        res = k - jump \n",
    "                        res = gg(stack[0], res)\n",
    "                        jump = stack[0+jump] - stack[0] \n",
    "                        v[x] = jump + res \n",
    "                    else:\n",
    "                        v[x] = stack[-1] - stack[len(stack)-k-2]\n",
    "                    \n",
    "                for nx in fa[x]:\n",
    "                    if nx not in v:\n",
    "                        dfs(nx, use) \n",
    "                stack.pop()\n",
    "            for j in dic:\n",
    "                dfs(j, 0)\n",
    "            # print(v)\n",
    "\n",
    "        \n",
    "        for j in c:\n",
    "            if j not in v:\n",
    "                f(j)\n",
    "\n",
    "        return max(v.values())\n",
    "class Solution:\n",
    "    def getMaxFunctionValue(self, r: List[int], k: int) -> int:\n",
    "        return mm(r, k)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def getMaxFunctionValue(self, receiver: List[int], k: int) -> int:\n",
    "        n = len(receiver)\n",
    "        parent = [[] for _ in range(n)]\n",
    "        for i in range(n):\n",
    "            x = receiver[i]\n",
    "            parent[x].append(i)\n",
    "\n",
    "        def findRoot():\n",
    "            root = []\n",
    "            visited = set()\n",
    "            for i in range(n):\n",
    "                loop = set()\n",
    "                while i not in visited:\n",
    "                    visited.add(i)\n",
    "                    loop.add(i)\n",
    "                    i = receiver[i]\n",
    "                if i in loop:\n",
    "                    root.append(i)\n",
    "            return root\n",
    "\n",
    "        class State:\n",
    "            def __init__(self, loop, k):\n",
    "                self.loop = loop\n",
    "                self.history = []\n",
    "\n",
    "                rep = k // len(self.loop)\n",
    "                rem = k % len(self.loop)\n",
    "                self.value = sum(self.loop) * rep\n",
    "                self.value += sum(self.loop[:rem])\n",
    "                self.loopIdx = (k - 1 + len(self.loop)) % len(self.loop)\n",
    "                self.loopK = k\n",
    "                self.hisIdx = -1\n",
    "\n",
    "            def push(self, x):\n",
    "                self.history.append(x)\n",
    "                saved = (self.loopIdx, self.loopK, self.hisIdx, self.value)\n",
    "\n",
    "                self.value += x\n",
    "                if self.hisIdx != -1:\n",
    "                    self.value -= self.history[self.hisIdx]\n",
    "                    self.hisIdx += 1\n",
    "                else:\n",
    "                    self.value -= self.loop[self.loopIdx]\n",
    "                    self.loopIdx -= 1\n",
    "                    self.loopK -= 1\n",
    "                    if self.loopIdx < 0 and self.loopK > 0:\n",
    "                        self.loopIdx += len(self.loop)\n",
    "                    if self.loopK == 0:\n",
    "                        self.hisIdx = 0\n",
    "                return saved\n",
    "\n",
    "            def pop(self, saved):\n",
    "                self.loopIdx, self.loopK, self.hisIdx, self.value = saved\n",
    "                self.history.pop()\n",
    "\n",
    "        root = findRoot()\n",
    "        ans = [-1] * n\n",
    "\n",
    "        def visitRoot(r):\n",
    "            loop = []\n",
    "            visit = set()\n",
    "            x = r\n",
    "            while x not in visit:\n",
    "                loop.append(x)\n",
    "                visit.add(x)\n",
    "                x = receiver[x]\n",
    "            st = State(loop, k + 1)\n",
    "\n",
    "            def dfs(r, st: State):\n",
    "                ans[r] = st.value\n",
    "                for p in parent[r]:\n",
    "                    if ans[p] == -1:\n",
    "                        saved = st.push(p)\n",
    "                        dfs(p, st)\n",
    "                        st.pop(saved)\n",
    "\n",
    "            dfs(r, st)\n",
    "\n",
    "        for r in root:\n",
    "            visitRoot(r)\n",
    "        # print(ans)\n",
    "        return max(ans)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def getMaxFunctionValue(self, receiver: List[int], k: int) -> int:\n",
    "        n = len(receiver)\n",
    "        parent = [[] for _ in range(n)]\n",
    "        for i in range(n):\n",
    "            x = receiver[i]\n",
    "            parent[x].append(i)\n",
    "\n",
    "        root = set()\n",
    "        visited = set()\n",
    "        for i in range(n):\n",
    "            loop = set()\n",
    "            while i not in visited:\n",
    "                visited.add(i)\n",
    "                loop.add(i)\n",
    "                i = receiver[i]\n",
    "            if i in loop:\n",
    "                root.add(i)\n",
    "\n",
    "        class State:\n",
    "            def __init__(self, loop, k):\n",
    "                self.loop = loop\n",
    "                self.history = []\n",
    "\n",
    "                rep = k // len(self.loop)\n",
    "                rem = k % len(self.loop)\n",
    "                self.value = sum(self.loop) * rep\n",
    "                self.value += sum(self.loop[:rem])\n",
    "\n",
    "                self.loopIdx = rem\n",
    "                self.loopK = k\n",
    "                self.hisIdx = -1\n",
    "\n",
    "            def push(self, x):\n",
    "                self.history.append(x)\n",
    "                saved = (self.loopIdx, self.loopK, self.hisIdx, self.value)\n",
    "\n",
    "                self.value += self.history[-1]\n",
    "                if self.hisIdx != -1:\n",
    "                    self.value -= self.history[self.hisIdx]\n",
    "                    self.hisIdx += 1\n",
    "                else:\n",
    "                    self.loopIdx -= 1\n",
    "                    self.loopK -= 1\n",
    "                    if self.loopK > 0 and self.loopIdx < 0:\n",
    "                        self.loopIdx += len(self.loop)\n",
    "                    if self.loopK == 0:\n",
    "                        self.hisIdx = 0\n",
    "                    self.value -= self.loop[self.loopIdx]\n",
    "                return saved\n",
    "\n",
    "\n",
    "            def pop(self, saved):\n",
    "                self.loopIdx, self.loopK, self.hisIdx, self.value = saved\n",
    "                self.history.pop()\n",
    "\n",
    "        # print(root)\n",
    "        ans = [-1] * n\n",
    "\n",
    "        def f(r):\n",
    "            loop = []\n",
    "            visit = set()\n",
    "            x = r\n",
    "            while x not in visit:\n",
    "                loop.append(x)\n",
    "                visit.add(x)\n",
    "                x = receiver[x]\n",
    "            st = State(loop, k+1)\n",
    "\n",
    "            def dfs(r, st: State):\n",
    "                ans[r] = st.value\n",
    "                for p in parent[r]:\n",
    "                    if ans[p] == -1:\n",
    "                        saved = st.push(p)\n",
    "                        dfs(p, st)\n",
    "                        st.pop(saved)\n",
    "\n",
    "            dfs(r, st)\n",
    "\n",
    "        for r in root:\n",
    "            f(r)\n",
    "\n",
    "        # print(ans)\n",
    "        return max(ans)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class UnionFind:\n",
    "    def __init__(self, length):\n",
    "        self.pa = [i for i in range(length)]\n",
    "\n",
    "    def find(self, i):\n",
    "        if self.pa[i] == i: return i\n",
    "        self.pa[i] = self.find(self.pa[i])\n",
    "        return self.pa[i]\n",
    "\n",
    "    def union(self, x, y):\n",
    "        self.pa[self.find(x)] = self.find(y)\n",
    "        \n",
    "class Solution:\n",
    "    def getMaxFunctionValue(self, receiver: List[int], k: int) -> int:\n",
    "        # 并查集 + 树上前缀和\n",
    "        uf = UnionFind(len(receiver))\n",
    "        n = len(receiver)\n",
    "        g = [[] for _ in range(n)]\n",
    "        for a, b in enumerate(receiver):\n",
    "            uf.union(a, b)\n",
    "            g[b].append(a)\n",
    "        roots = set(uf.find(i) for i in uf.pa)\n",
    "        ps = {}\n",
    "        for r in roots:\n",
    "            o = receiver[r]\n",
    "            ps[r] = [r]\n",
    "            while o != r:\n",
    "                ps[r].append(ps[r][-1] + o)\n",
    "                o = receiver[o]\n",
    "\n",
    "        res = [-1] * n\n",
    "        for r in roots:\n",
    "            treesum = [0, r]\n",
    "            vis = {r}\n",
    "            def dfs(u):\n",
    "                tk = k + 2 - len(treesum)\n",
    "                if tk <= 0:\n",
    "                    res[u] = treesum[-1] - treesum[- 2 - k]\n",
    "                else:\n",
    "                    nc = tk // len(ps[r])\n",
    "                    rem = tk % len(ps[r])\n",
    "                    res[u] = treesum[-1] + nc * ps[r][-1] + ps[r][rem] - r\n",
    "\n",
    "                for v in g[u]:\n",
    "                    if v in vis: continue\n",
    "                    vis.add(v)\n",
    "                    treesum.append(treesum[-1] + v)\n",
    "                    dfs(v)\n",
    "                    treesum.pop()\n",
    "            dfs(r)\n",
    "        return max(res)\n",
    "                \n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def getMaxFunctionValue(self, receiver: List[int], k: int) -> int:\n",
    "        n=len(receiver)\n",
    "        M=34\n",
    "        dp=[[0]*n for _ in range(M)]\n",
    "        toreceiverid=[[0]*n for _ in range(M)]\n",
    "        for i in range(n):\n",
    "            dp[0][i]=i+receiver[i]\n",
    "            toreceiverid[0][i]=receiver[i]\n",
    "        for i in range(1,M):\n",
    "            for j in range(n):\n",
    "                dp[i][j]=dp[i-1][j]+dp[i-1][toreceiverid[i-1][j]]-toreceiverid[i-1][j]\n",
    "                toreceiverid[i][j]=toreceiverid[i-1][toreceiverid[i-1][j]]\n",
    "        res=0\n",
    "        for i in range(n):\n",
    "            tmp=i\n",
    "            nowid=i\n",
    "            for i in range(M):\n",
    "                if ((1<<i)&k) !=0:\n",
    "                    tmp=tmp+dp[i][nowid]-nowid\n",
    "                    nowid=toreceiverid[i][nowid]\n",
    "                res=max(res,tmp)\n",
    "        return res\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def getMaxFunctionValue(self, r: List[int], k: int) -> int:\n",
    "        \"\"\"\n",
    "            选择 x\n",
    "            x -> r[x] -> r[r[x]] -> ...\n",
    "            是否形成循环?\n",
    "            \n",
    "            走k步 构成\n",
    "            循环内+循环外\n",
    "        \"\"\"\n",
    "        n = len(r)\n",
    "        \n",
    "        # 最多走2^l步\n",
    "        l = 0\n",
    "        tmp = 1\n",
    "        while tmp <= k:\n",
    "            tmp *= 2\n",
    "            l += 1\n",
    "        \n",
    "        rl = [[0] * n for _ in range(l + 1)]\n",
    "        for i in range(n):\n",
    "            rl[0][i] = r[i] # 走1步\n",
    "            # rl[1][i] = r[i]\n",
    "        \n",
    "        for step in range(1, l + 1):\n",
    "            for i in range(n):\n",
    "                rl[step][i] = rl[step - 1][rl[step - 1][i]]\n",
    "        \n",
    "        # 从 i位置走了 (1<<j) 步 到了那个下标\n",
    "        # rl[j][i]的意义\n",
    "        # print(rl)\n",
    "        \n",
    "        # sm[j][i]: 从i位置走了 (1<<j) 步 所得的总和\n",
    "        sm = [[0] * n for _ in range(l + 1)]\n",
    "        for i in range(n):\n",
    "            sm[0][i] = i + r[i] # 走(1 << 0)步\n",
    "            # sm[1][i] = i + r[i]\n",
    "        \n",
    "        for step in range(1, l + 1):\n",
    "            for i in range(n):\n",
    "                sm[step][i] = sm[step - 1][rl[step - 1][i]] + sm[step - 1][i] - rl[step - 1][i] # 中间多走1次\n",
    "        # print(sm)\n",
    "        \n",
    "        ans = 0\n",
    "        for i in range(n):\n",
    "            res = 0\n",
    "            kk = k\n",
    "            idx = i\n",
    "            for bit in range(l, -1, -1):\n",
    "                p2 = pow(2, bit)\n",
    "                if kk >= p2: \n",
    "                    # print(\"从{} 走了 {}步, 产生了{}的代价\".format(idx, (1 << bit), sm[bit][idx]))\n",
    "                    res += sm[bit][idx]\n",
    "                    kk -= p2\n",
    "                    # 从idx开始走了(1<<bit)步的新位置\n",
    "                    idx = rl[bit][idx]\n",
    "                    if kk > 0: # 除了最后一步\n",
    "                        res -= idx\n",
    "            ans = max(ans, res)\n",
    "        \n",
    "        # cao[i][j]: i位置元素走了1<<j步的和大小为多少\n",
    "        \n",
    "        # cao[i][2j] = cao[i][j] + cao[rj[i]][j]\n",
    "        \"\"\"\n",
    "            i -> r[i] -> r2[i]\n",
    "            cao[i][2] = cao[r[i]][1]\n",
    "            cao[i][4] = cao[r[r[i]]][2]\n",
    "        \n",
    "        \"\"\"\n",
    "        \n",
    "            \n",
    "        return ans\n",
    "            \n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def getMaxFunctionValue(self, receiver: List[int], k: int) -> int:\n",
    "        n=len(receiver)\n",
    "        dp=[[0]*n for _ in range(36)]\n",
    "        toreceiverid=[[0]*n for _ in range(36)]\n",
    "        for i in range(n):\n",
    "            dp[0][i]=i+receiver[i]\n",
    "            toreceiverid[0][i]=receiver[i]\n",
    "        for i in range(1,36):\n",
    "            for j in range(n):\n",
    "                dp[i][j]=dp[i-1][j]+dp[i-1][toreceiverid[i-1][j]]-toreceiverid[i-1][j]\n",
    "                toreceiverid[i][j]=toreceiverid[i-1][toreceiverid[i-1][j]]\n",
    "        res=0\n",
    "        for i in range(n):\n",
    "            tmp=i\n",
    "            nowid=i\n",
    "            for i in range(36):\n",
    "                if ((1<<i)&k) !=0:\n",
    "                    tmp=tmp+dp[i][nowid]-nowid\n",
    "                    nowid=toreceiverid[i][nowid]\n",
    "                res=max(res,tmp)\n",
    "        return res\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def getMaxFunctionValue(self, receiver: List[int], k: int) -> int:\n",
    "        n = len(receiver)\n",
    "        tree = [[0]*30 for _ in range(n)]\n",
    "        parent = [[-1]*30 for _ in range(n)]\n",
    "        ps = [1]*30\n",
    "            \n",
    "        for i in range(n):\n",
    "            tree[i][0] = receiver[i]  \n",
    "            parent[i][0] = receiver[i]  \n",
    "        for p in range(1,30):\n",
    "            ps[p] = 2*ps[p-1]\n",
    "            for i in range(n):\n",
    "                parent[i][p] = parent[parent[i][p-1]][p-1]\n",
    "                tree[i][p] = tree[i][p-1]+tree[parent[i][p-1]][p-1]\n",
    "\n",
    "        # print(tree)\n",
    "        # print(parent)\n",
    "        # print(ps)\n",
    "\n",
    "        ans = 0\n",
    "        for i in range(n):\n",
    "            s = k \n",
    "            j = 29\n",
    "            cur = i\n",
    "            while s:\n",
    "                while j >= 0 and ps[j] > s:\n",
    "                    j -= 1\n",
    "                s -= ps[j]\n",
    "                cur += tree[i][j]\n",
    "                # print(i,j,parent[i][j])\n",
    "                i = parent[i][j]\n",
    "            # print(i,cur)\n",
    "            ans = max(ans,cur)\n",
    "        return ans \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def getMaxFunctionValue(self, receiver: List[int], k: int) -> int:\n",
    "        n=len(receiver)\n",
    "        dp=[[0]*36 for _ in range(n)]\n",
    "        toreceiverid=[[0]*36 for _ in range(n)]\n",
    "        for i in range(n):\n",
    "            dp[i][0]=i+receiver[i]\n",
    "            toreceiverid[i][0]=receiver[i]\n",
    "        for i in range(1,36):\n",
    "            for j in range(n):\n",
    "                dp[j][i]=dp[j][i-1]+dp[toreceiverid[j][i-1]][i-1]-toreceiverid[j][i-1]\n",
    "                toreceiverid[j][i]=toreceiverid[toreceiverid[j][i-1]][i-1]\n",
    "        res=0\n",
    "        for i in range(n):\n",
    "            tmp=i\n",
    "            nowid=i\n",
    "            for i in range(36):\n",
    "                if ((1<<i)&k) !=0:\n",
    "                    tmp=tmp+dp[nowid][i]-nowid\n",
    "                    nowid=toreceiverid[nowid][i]\n",
    "                res=max(res,tmp)\n",
    "        return res\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def getMaxFunctionValue(self, receiver: List[int], k: int) -> int:\n",
    "        n = len(receiver)\n",
    "        m = math.floor(log(k, 2))\n",
    "        f = [[0] * n for _ in range(m+1)]\n",
    "        g = [[-1] * n for _ in range(m+1)]\n",
    "        for i in range(n):\n",
    "            f[0][i] = g[0][i] = receiver[i]\n",
    "        \n",
    "        for j in range(1, m+1):\n",
    "            for i in range(n):\n",
    "                f[j][i] = f[j-1][f[j-1][i]]\n",
    "                g[j][i] = g[j-1][i] + g[j-1][f[j-1][i]]\n",
    "\n",
    "        res = 0\n",
    "        for i in range(n): \n",
    "            v = k\n",
    "            j = m \n",
    "            p = t = i\n",
    "            while v > 0:\n",
    "                while 1<< j > v:\n",
    "                    j -= 1\n",
    "                t += g[j][p]\n",
    "                v -= 1<<j \n",
    "                p = f[j][p]\n",
    "            res = max(res, t)\n",
    "        return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def getMaxFunctionValue(self, receiver: List[int], k: int) -> int:\n",
    "        n = len(receiver)\n",
    "        logn = k.bit_length()\n",
    "        # 从第 j 个人开始传递 2^i 次后的总得分 (从第 j 个人开始一共计算 2^i 次分数)\n",
    "        dp = [[0] * n for _ in range(logn)]\n",
    "        # 第 j 个人传递 2^i 次后球的位置\n",
    "        pos = [[0] * n for _ in range(logn)]\n",
    "        for i in range(n):\n",
    "            dp[0][i] = receiver[i]\n",
    "            pos[0][i] = receiver[i]\n",
    "\n",
    "        for i in range(1, logn):\n",
    "            for j in range(n):\n",
    "                p = pos[i - 1][j]\n",
    "                dp[i][j] = dp[i - 1][j] + dp[i - 1][p]\n",
    "                pos[i][j] = pos[i - 1][p]\n",
    "\n",
    "        ans = list(range(n))\n",
    "        np = list(range(n))\n",
    "        for i in range(logn):\n",
    "            if k & (1 << i):\n",
    "                for j in range(n):\n",
    "                    ans[j] += dp[i][np[j]]\n",
    "                for j in range(n):\n",
    "                    np[j] = pos[i][np[j]]\n",
    "\n",
    "        return max(ans)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def getMaxFunctionValue(self, receiver: List[int], k: int) -> int:\n",
    "        n, N = len(receiver), 34\n",
    "        f = [[0] * N for _ in range(n)]\n",
    "        g = [[0] * N for _ in range(n)]\n",
    "        # f[i, p] = f[i, p - 1] + f[g[i, p - 1], p - 1]\n",
    "        # g[i, p] = g[g[i, p - 1], p - 1]\n",
    "        \n",
    "        for i in range(n):\n",
    "            f[i][0], g[i][0] = receiver[i], receiver[i]\n",
    "\n",
    "        for j in range(1, N):\n",
    "            for i in range(n):\n",
    "                f[i][j] = f[i][j - 1] + f[g[i][j - 1]][j - 1]\n",
    "                g[i][j] = g[g[i][j - 1]][j - 1]\n",
    "        ans = 0\n",
    "        for i in range(n):\n",
    "            res = t = i\n",
    "            for j in range(N):\n",
    "                if k >> j & 1:\n",
    "                    res += f[t][j]\n",
    "                    t = g[t][j]\n",
    "            ans = max(ans, res)\n",
    "        return ans\n",
    "                    \n",
    "                    \n",
    "            \n",
    "                                \n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "#\n",
    "# @lc app=leetcode.cn id=2836 lang=python\n",
    "#\n",
    "# [2836] 在传球游戏中最大化函数值\n",
    "#\n",
    "\n",
    "# @lc code=start\n",
    "class Solution(object):\n",
    "    def getMaxFunctionValue(self, receiver, k):\n",
    "        \"\"\"\n",
    "        :type receiver: List[int]\n",
    "        :type k: int\n",
    "        :rtype: int\n",
    "        \"\"\"\n",
    "        target = [[0] * 36 for i in range(len(receiver))]\n",
    "        path_sum = [[0] * 36 for i in range(len(receiver))]\n",
    "        time = 0\n",
    "        copy_k = k\n",
    "        while copy_k > 0:\n",
    "            time += 1\n",
    "            copy_k = copy_k // 2\n",
    "        max_val = 0\n",
    "        for i in range(len(receiver)):\n",
    "            target[i][0] = receiver[i]\n",
    "            path_sum[i][0] = receiver[i]\n",
    "            max_val = max(max_val, path_sum[i][0])\n",
    "\n",
    "        for i in range(1,time):\n",
    "            for j in range(len(receiver)):\n",
    "                target[j][i] = target[target[j][i - 1]][i-1]\n",
    "                path_sum[j][i] = path_sum[j][i - 1] + path_sum[target[j][i - 1]][i-1]\n",
    "        #print(path_sum)\n",
    "        for i in range(len(receiver)):\n",
    "            copy_k = k\n",
    "            now_target = i\n",
    "            now_path = 0\n",
    "            count = 0\n",
    "            while copy_k> 0:\n",
    "                if copy_k % 2 == 1:\n",
    "                    now_path += path_sum[now_target][count]\n",
    "                    now_target = target[now_target][count]\n",
    "                count += 1\n",
    "                copy_k = copy_k //2\n",
    "            max_val = max(max_val, now_path + i)\n",
    "\n",
    "        return max_val\n",
    "\n",
    "# @lc code=end\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def getMaxFunctionValue(self, rec: List[int], k: int) -> int:\n",
    "        n = len(rec)\n",
    "        m = k.bit_length()\n",
    "        f = [[p] + [0] * m for p in rec]\n",
    "        v = [[p] + [0] * m for p in rec]\n",
    "\n",
    "        for i in range(m):\n",
    "            for x in range(n):\n",
    "                pa = f[x][i]\n",
    "                f[x][i + 1] = f[pa][i]\n",
    "                v[x][i + 1] = v[x][i] + v[pa][i]\n",
    "\n",
    "        ans = 0\n",
    "        for x in range(n):\n",
    "            s = x\n",
    "            for i in range(m):\n",
    "                if (k >> i) & 1 != 0:\n",
    "                    s += v[x][i]\n",
    "                    x = f[x][i]\n",
    "            ans = max(ans, s)\n",
    "\n",
    "        return ans\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def getMaxFunctionValue(self, receiver: List[int], k: int) -> int:\n",
    "\n",
    "        n = len(receiver)\n",
    "        dp = [[0]*36 for _ in range(n)]\n",
    "        weight = [[0]*36 for _ in range(n)]\n",
    "        for i in range(n):\n",
    "            dp[i][0] = receiver[i]\n",
    "            weight[i][0] = i\n",
    "        \n",
    "        for j in range(1,36):\n",
    "            for i in range(n):\n",
    "                dp[i][j] = dp[dp[i][j-1]][j-1]\n",
    "                weight[i][j] = weight[i][j-1]+weight[dp[i][j-1]][j-1]\n",
    "        \n",
    "        res = 0 \n",
    "\n",
    "        for i in range(n):\n",
    "            cur = 0\n",
    "            pos = i \n",
    "            for j in range(36):\n",
    "                if (k>>j)&1:\n",
    "                    cur += weight[pos][j]\n",
    "                    pos = dp[pos][j]\n",
    "            res = max(res, cur+pos)\n",
    "        return res\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def getMaxFunctionValue(self, receiver: List[int], k: int) -> int:\n",
    "        n = len(receiver)\n",
    "        to = [[0]*36 for _ in range(n)]\n",
    "        cost = [[0]*36 for _ in range(n)]\n",
    "        for i in range(n):\n",
    "            to[i][0] = receiver[i]\n",
    "            cost[i][0] = i\n",
    "        \n",
    "        for j in range(1,36):\n",
    "            for i in range(n):\n",
    "                 to[i][j] = to[to[i][j-1]][j-1]\n",
    "                 cost[i][j] = cost[i][j-1]+cost[to[i][j-1]][j-1]\n",
    "\n",
    "        res = 0\n",
    "        for i in range(n):\n",
    "            cur = 0\n",
    "            pos = i\n",
    "            for j in range(36):\n",
    "                if (k>>j & 1):\n",
    "                    cur += cost[pos][j]\n",
    "                    pos = to[pos][j]\n",
    "            res = max(res,cur+pos)\n",
    "        return res\n",
    "        \n",
    "\n",
    "\n",
    "        \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "MX, LY = 100000, 35\n",
    "ps = [[0 for _ in range(LY)] for _ in range(MX)]\n",
    "gold = [[0 for _ in range(LY)] for _ in range(MX)]\n",
    "\n",
    "class Solution:\n",
    "    def getMaxFunctionValue(self, rcv: List[int], k: int) -> int:\n",
    "        for i in range(len(rcv)):\n",
    "            ps[i][0] = rcv[i]\n",
    "            gold[i][0] = i\n",
    "        \n",
    "        for i in range(1, LY):\n",
    "            for j in range(len(rcv)):\n",
    "                ps[j][i] = ps[ps[j][i-1]][i-1]\n",
    "                gold[j][i] = gold[ps[j][i-1]][i-1] + gold[j][i-1]\n",
    "\n",
    "        ans = 0\n",
    "        for i in range(len(rcv)):\n",
    "            sum = 0\n",
    "            r = i\n",
    "            for j in range(LY):\n",
    "                if k & 1<<j:\n",
    "                    sum += gold[r][j]\n",
    "                    r = ps[r][j]\n",
    "                   \n",
    "            sum += r\n",
    "            ans = max(sum, ans)\n",
    "\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "def binary(num):\n",
    "    ls=[]\n",
    "    x=num\n",
    "    while x>0:\n",
    "        b=x%2\n",
    "        x=x//2\n",
    "        ls.append(b)\n",
    "    ls=ls+[0]*(40-len(ls))\n",
    "    return ls\n",
    "class Solution:\n",
    "    def getMaxFunctionValue(self, receiver: List[int], k: int) -> int:\n",
    "        n=len(receiver)\n",
    "        dpls=[[0 for i in range(n)] for j in range(40)]\n",
    "        dpls[0]=receiver\n",
    "        wls=[[0 for i in range(n)] for j in range(40)]\n",
    "        wls[0]=[i for i in range(n)]\n",
    "        for i in range(1,40):\n",
    "            for j in range(n):\n",
    "                dpls[i][j]=dpls[i-1][dpls[i-1][j]]\n",
    "                wls[i][j]=wls[i-1][j]+wls[i-1][dpls[i-1][j]]\n",
    "        kls=binary(k)\n",
    "        rls=[0 for i in range(n)]\n",
    "        for j in range(n):\n",
    "            index=j\n",
    "            for i in range(40):\n",
    "                if kls[i]==1:\n",
    "                    rls[j]+=wls[i][index]\n",
    "                    index=dpls[i][index]\n",
    "            rls[j]+=index\n",
    "        return max(rls)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def getMaxFunctionValue(self, receiver: List[int], k: int) -> int:\n",
    "        n = len(receiver)\n",
    "        f = [[0] * 36 for _ in range(n)]\n",
    "        w = [[0] * 36 for _ in range(n)]\n",
    "        for i in range(n):\n",
    "            f[i][0] = receiver[i]\n",
    "            w[i][0] = i \n",
    "        for j in range(1, 36):\n",
    "            for i in range(n):\n",
    "                f[i][j] = f[f[i][j-1]][j-1]\n",
    "                w[i][j] = w[i][j-1] + w[f[i][j-1]][j-1] \n",
    "        @cache\n",
    "        def calc(x, k):\n",
    "            if k == 0:\n",
    "                return x \n",
    "            tot = 0 \n",
    "            pos = x \n",
    "            for j in range(36):\n",
    "                if (k >> j) & 1:\n",
    "                    tot += w[pos][j]\n",
    "                    pos = f[pos][j]\n",
    "            return tot + pos \n",
    "\n",
    "        ans = 0 \n",
    "        for i in range(n):\n",
    "            ans = max(ans, calc(i, k))\n",
    "        return ans         \n",
    "                    \n",
    "            \n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def getMaxFunctionValue(self, receiver: List[int], m: int) -> int:\n",
    "        n = len(receiver)\n",
    "        f = [[i] + [0] * 40 for i in range(n)]\n",
    "        p = [[receiver[i]] + [0] * 40 for i in range(n)]\n",
    "        for k in range(40):\n",
    "            for i in range(n):\n",
    "                f[i][k + 1] = f[i][k] + f[p[i][k]][k]\n",
    "                p[i][k + 1] = p[p[i][k]][k]\n",
    "        \n",
    "        m += 1\n",
    "        ans = 0\n",
    "        for i in range(n):\n",
    "            cur = 0\n",
    "            for k in range(40, -1, -1):\n",
    "                if m >> k & 1:\n",
    "                    cur += f[i][k]\n",
    "                    i = p[i][k]\n",
    "            ans = max(ans, cur)\n",
    "        \n",
    "        return ans\n",
    "                \n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def getMaxFunctionValue(self, receiver: List[int], kk: int) -> int:\n",
    "        n = len(receiver)\n",
    "        ind = [0]*n\n",
    "        #kk = kk+1\n",
    "        for i,a in enumerate(receiver):\n",
    "            ind[a] +=1\n",
    "        dic = defaultdict(list)\n",
    "        def dfs(i,acc,idx,k):\n",
    "            visit2[i] =1\n",
    "            acc += i \n",
    "            if len(dic[i]) ==2:\n",
    "                return\n",
    "            if len(dic[i])>0 and dic[i][0][2] != k:\n",
    "                return\n",
    "            dic[i].append((acc,idx,k))\n",
    "            dfs(receiver[i],acc,idx+1,k)\n",
    "        rings = {}\n",
    "        rrings = defaultdict(list)\n",
    "        visit2 ={}\n",
    "        for i in range(n):\n",
    "            if ind[i] ==0:\n",
    "                dfs(i,0,1,i)\n",
    "            #print(dic,i)\n",
    "        for i in range(n):\n",
    "            if i not in visit2:\n",
    "                dfs(i,0,1,i)\n",
    "\n",
    "        for i in range(n):\n",
    "            if len(dic[i]) ==2:\n",
    "                #print(dic[i])\n",
    "                rings[i] = dic[i][1][2]\n",
    "                rrings[dic[i][1][2]].append(i)\n",
    "        #print(rings,rrings)\n",
    "        #print(dic)\n",
    "        visit = {}\n",
    "        mx = 0\n",
    "        acc = 0\n",
    "        for idx in range(n):\n",
    "            if ind[idx] ==0:\n",
    "                acc=idx\n",
    "                i = idx\n",
    "                k=kk\n",
    "                #print(idx,k,i,acc)\n",
    "                while k>0:\n",
    "                    k -=1\n",
    "                    i = receiver[i]\n",
    "                    acc +=i\n",
    "                    #print(idx,k,i,acc,rings)\n",
    "                    if i in rings and k> len(rrings[rings[i]]):\n",
    "                        t = k // len(rrings[rings[i]])\n",
    "                        k = k % len(rrings[rings[i]])\n",
    "                        acc += (dic[i][1][0] -dic[i][0][0])*t\n",
    "                #print(idx,acc)\n",
    "                while idx not in visit:\n",
    "                    mx  =max(mx,acc)\n",
    "                    visit[idx] =1 \n",
    "                    acc -=idx\n",
    "                    idx = receiver[idx]\n",
    "                    i = receiver[i]\n",
    "                    acc += i\n",
    "        for idx in range(n):\n",
    "            if idx not in visit:\n",
    "                acc=0\n",
    "                acc += idx\n",
    "                i = idx\n",
    "                k=kk\n",
    "                while k>0:\n",
    "                    k -=1\n",
    "                    i = receiver[i]\n",
    "                    acc +=i\n",
    "                    #print(idx,k,i,acc)\n",
    "                    if i in rings and k> len(rrings[rings[i]]):\n",
    "                        t = k // len(rrings[rings[i]])\n",
    "                        k = k % len(rrings[rings[i]])\n",
    "                        acc += (dic[i][1][0] -dic[i][0][0])*t\n",
    "                while idx not in visit:\n",
    "                    #print(idx,acc,i)\n",
    "                    mx  =max(mx,acc)\n",
    "                    visit[idx] =1 \n",
    "                    acc -=idx\n",
    "                    idx = receiver[idx]\n",
    "                    i = receiver[i]\n",
    "\n",
    "                    acc += i\n",
    "        #print(visit)\n",
    "        return mx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "import numpy as np\n",
    "class Solution:\n",
    "    def getMaxFunctionValue(self, receiver: List[int], k: int) -> int:\n",
    "        N = len(receiver)\n",
    "        nxt = [[j for j in range(50)] for i in range(N+1)] \n",
    "        val = [[j for j in range(50)] for i in range(N+1)] \n",
    "        \n",
    "        for i in range(N):\n",
    "            nxt[i][0] = receiver[i]\n",
    "            val[i][0] = i\n",
    "        \n",
    "        for j in range(40):\n",
    "            for i in range(N):\n",
    "                nxt[i][j + 1] = nxt[nxt[i][j]][j]\n",
    "                val[i][j + 1] = val[i][j] + val[nxt[i][j]][j]\n",
    "        \n",
    "        ans = 0\n",
    "\n",
    "        for i in range(N):\n",
    "            tmp = 0\n",
    "            start = i\n",
    "            for j in range(40):\n",
    "                if ((k>>j) & 1):\n",
    "                    tmp += val[start][j]\n",
    "                    start = nxt[start][j]\n",
    "            ans = max(ans, tmp + start)\n",
    "        \n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "import numpy as np\n",
    "class Solution:\n",
    "    def getMaxFunctionValue(self, receiver: List[int], k: int) -> int:\n",
    "        N = len(receiver)\n",
    "        nxt = [[j for j in range(50)] for i in range(N+1)] \n",
    "        val = [[j for j in range(50)] for i in range(N+1)] \n",
    "        \n",
    "        for i in range(N):\n",
    "            nxt[i][0] = receiver[i]\n",
    "            val[i][0] = i\n",
    "        \n",
    "        for j in range(45):\n",
    "            for i in range(N):\n",
    "                nxt[i][j + 1] = nxt[nxt[i][j]][j]\n",
    "                val[i][j + 1] = val[i][j] + val[nxt[i][j]][j]\n",
    "        \n",
    "        ans = 0\n",
    "\n",
    "        for i in range(N):\n",
    "            tmp = 0\n",
    "            start = i\n",
    "            for j in range(45):\n",
    "                if ((k>>j) & 1):\n",
    "                    tmp += val[start][j]\n",
    "                    start = nxt[start][j]\n",
    "            ans = max(ans, tmp + start)\n",
    "        \n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def getMaxFunctionValue(self, receiver: List[int], k: int) -> int:\n",
    "        n = len(receiver)\n",
    "        dp = [[0 for __ in range(50)] for _ in range(n)]\n",
    "        dis = [[0 for __ in range(50)] for _ in range(n)]\n",
    "        for i in range(n):\n",
    "            dp[i][0] = receiver[i]\n",
    "            dis[i][0] = i\n",
    "        for i in range(1, 50):\n",
    "            for j in range(n):\n",
    "                mid = dp[j][i-1]\n",
    "                d1 = dis[j][i-1]\n",
    "                t = dp[mid][i-1]\n",
    "                d2 = dis[mid][i-1]\n",
    "                dp[j][i] = t\n",
    "                dis[j][i] = d1+d2\n",
    "        ans = 0\n",
    "        for i in range(n):\n",
    "            st = i\n",
    "            tmp = 0\n",
    "            for j in range(39, -1, -1):\n",
    "                if k & (1 << j):\n",
    "                    tmp += dis[st][j]\n",
    "                    st = dp[st][j]\n",
    "            ans = max(ans, tmp+st)\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def getMaxFunctionValue(self, receiver: List[int], k: int) -> int:\n",
    "        N = len(receiver)\n",
    "        LOG = int(math.log(k,2)) + 1\n",
    "\n",
    "        cache = [[(-114514, -114514) for __ in range(LOG)] for _ in range(N)]\n",
    "        def ancestor(n, j):\n",
    "            if j == 0:\n",
    "                return receiver[n], n + receiver[n]\n",
    "            if cache[n][j][0] != -114514:\n",
    "                return cache[n][j]\n",
    "            half_way_node, cost = ancestor(n, j - 1)\n",
    "            final, cost_2 = ancestor(half_way_node, j - 1)\n",
    "            res = final, cost + cost_2 - half_way_node\n",
    "            cache[n][j] = res\n",
    "            return res\n",
    "        def getKthAncestor(node: int, k: int) -> int:\n",
    "            cost = 0\n",
    "            prev_node_val = 0\n",
    "            for i in range(LOG):\n",
    "                if k & (1 << i) != 0:\n",
    "                    node, dc = ancestor(node, i)\n",
    "                    cost += dc - prev_node_val\n",
    "                    prev_node_val = node\n",
    "            return cost\n",
    "        res = 0\n",
    "        for i in range((len(receiver))):\n",
    "            res = max(res, getKthAncestor(i, k))\n",
    "        return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def getMaxFunctionValue(self, receiver: List[int], k: int) -> int:\n",
    "        n = len(receiver)\n",
    "        mx = k.bit_length()\n",
    "        F = [[None] * n for i in range(mx)]\n",
    "        F[0] = [(i,i) for i in receiver]\n",
    "        for i in range(1,mx):\n",
    "            for strt in range(n):\n",
    "                end,v = F[i-1][strt]\n",
    "                F[i][strt] = (F[i-1][end][0],v+F[i-1][end][1])\n",
    "        res = -inf\n",
    "        for i in range(n):\n",
    "            temp = i;x = k;pos = 0;strt = i\n",
    "            while x:\n",
    "                if x & 1:\n",
    "                    strt,add = F[pos][strt]\n",
    "                    temp += add\n",
    "                pos += 1\n",
    "                x >>= 1\n",
    "            res = max(res,temp)\n",
    "        return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def getMaxFunctionValue(self, g: List[int], k: int) -> int:\n",
    "        \n",
    "        n = len(g)\n",
    "        m = k.bit_length()-1\n",
    "        fa = [[(p,p)] + [None]*m for p in g]\n",
    "\n",
    "        for j in range(m):\n",
    "            for i in range(n):\n",
    "                p,s = fa[i][j]\n",
    "                pp, ss = fa[p][j]\n",
    "                fa[i][j + 1] = (pp, s + ss)  # 合并节点值之和\n",
    "        \n",
    "        res = 0\n",
    "        for i in range(n):\n",
    "            x = sums = i\n",
    "            for j in range(m+1):\n",
    "                if (k>>j)&1:\n",
    "                    x,s = fa[x][j]\n",
    "                    sums+=s\n",
    "            res = max(res,sums)\n",
    "        \n",
    "        return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def getMaxFunctionValue(self, receiver: List[int], k: int) -> int:\n",
    "        n = len(receiver)\n",
    "        m = k.bit_length() - 1\n",
    "        pa = [[(p, p)] + [None] * m for p in receiver]\n",
    "        for i in range(m):\n",
    "            for x in range(n):\n",
    "                p, s = pa[x][i]\n",
    "                pp, ss = pa[p][i]\n",
    "                pa[x][i + 1] = (pp, s + ss)  # 合并节点值之和\n",
    "\n",
    "        ans = 0\n",
    "        for i in range(n):\n",
    "            x = sum = i\n",
    "            for j in range(m + 1):\n",
    "                if (k >> j) & 1:  # k 的二进制从低到高第 j 位是 1\n",
    "                    x, s = pa[x][j]\n",
    "                    sum += s\n",
    "            ans = max(ans, sum)\n",
    "        return ans\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def getMaxFunctionValue(self, receiver, k: int) -> int:\n",
    "        sz = len(receiver)\n",
    "        bits = []\n",
    "        for x in range(64):\n",
    "            if (1 << x) & (k + 1):\n",
    "                bits.append(x)\n",
    "\n",
    "        d = [[0 for y in range(bits[-1] + 1)] for x in range(sz)]\n",
    "        for x in range(sz):\n",
    "            d[x][0] = (receiver[x], x)\n",
    "\n",
    "        for j in range(1, bits[-1] + 1):\n",
    "            for i in range(sz):\n",
    "                d[i][j] = (d[d[i][j - 1][0]][j - 1][0], d[i][j - 1][1] + d[d[i][j - 1][0]][j - 1][1])\n",
    "\n",
    "        ret = 0\n",
    "        for y in range(sz):\n",
    "            t = 0\n",
    "            nxt = y\n",
    "            for x in bits:\n",
    "                t += d[nxt][x][1]\n",
    "                nxt = d[nxt][x][0]\n",
    "            ret = max(ret, t)\n",
    "        return ret"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def getMaxFunctionValue(self, receiver: List[int], k: int) -> int:\n",
    "        m = k.bit_length()\n",
    "        n = len(receiver)\n",
    "        pa = [[(p,p)] + [None] * m for p in receiver]\n",
    "        for i in range(m):\n",
    "            for x in range(n):\n",
    "                p, s = pa[x][i]\n",
    "                pp, ss = pa[p][i]\n",
    "                pa[x][i + 1] = (pp, s + ss)\n",
    "        ans = 0\n",
    "        for i in range(n):\n",
    "            x = sum = i\n",
    "            for j in range(m + 1):\n",
    "                if k >> j & 1:\n",
    "                    x, s = pa[x][j]\n",
    "                    sum += s\n",
    "            ans = max(ans, sum)\n",
    "        return ans\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def getMaxFunctionValue(self, receiver: List[int], k: int) -> int:\n",
    "        n = len(receiver)\n",
    "        m = k.bit_length()\n",
    "        f = [[(v, i)]+[(-1, 0)]*m for i, v in enumerate(receiver)]\n",
    "        for j in range(m):\n",
    "            for i in range(n):\n",
    "                a, b = f[i][j]\n",
    "                a2, b2 = f[a][j]\n",
    "                f[i][j+1] = (a2, b2+b)\n",
    "        ret = 0\n",
    "        for i in range(n):\n",
    "            x = receiver[i]\n",
    "            t = i\n",
    "            for i in range(k.bit_length()):\n",
    "                if (k >> i) & 1:\n",
    "                    t += f[x][i][1]\n",
    "                    x = f[x][i][0]\n",
    "            ret = max(ret, t)\n",
    "        return ret"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def getMaxFunctionValue(self, receiver: List[int], k: int) -> int:\n",
    "        m = k.bit_length()\n",
    "        n = len(receiver)\n",
    "        dp = [[[0, 0] for _ in range(n)] for _ in range(m)]\n",
    "        dp[0] = [[v, x + v] for x, v in enumerate(receiver)]\n",
    "        for i in range(m - 1):\n",
    "            for j in range(n):\n",
    "                p, s = dp[i][j]\n",
    "                pp, ss = dp[i][p]\n",
    "                dp[i + 1][j] = [pp, s + ss - p]\n",
    "        ans = 0\n",
    "        for x in range(n):\n",
    "            res = 0\n",
    "            for i in range(m - 1, -1, -1):\n",
    "                if (k >> i) & 1:\n",
    "                    p, s = dp[i][x]\n",
    "                    if i == m - 1:\n",
    "                        res += s\n",
    "                    else:\n",
    "                        res += (s - x)\n",
    "                    x = p\n",
    "            ans = max(ans, res)\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class TreeDoubling:\n",
    "    '''\n",
    "    parents: 父节点关系的数组，parents[i] 是 i 的父节点\n",
    "    k: 最大查询范围，可以很大（10**10）,反正会指数折叠\n",
    "    '''\n",
    "    def __init__(self, parents, k):\n",
    "        m = k.bit_length()\n",
    "        pa = {}\n",
    "        for x, p in parents.items():\n",
    "            y = 0 if p is None else p\n",
    "            pa[x] = [(p, y)] + [(None, 0)] * m\n",
    "        pa[None] = [(None, 0)] * m\n",
    "        for i in range(m):\n",
    "            for x in parents:\n",
    "                p, s = pa[x][i]\n",
    "                pp, ps = pa[p][i]\n",
    "                pa[x][i+1] = (pp, ps + s)\n",
    "        self.pa = pa\n",
    "\n",
    "    def get_Kth_ancestor(self, node, k):\n",
    "        s0 = node\n",
    "        for i in range(k.bit_length()):\n",
    "            if (k >> i) & 1:\n",
    "                node, s = self.pa[node][i]\n",
    "                if node is None: return None, 0\n",
    "                s0 += s # 与 ps + s 的累加关系对应修改\n",
    "        return node, s0\n",
    "\n",
    "\n",
    "class Solution:\n",
    "    def getMaxFunctionValue(self, receiver: List[int], k: int) -> int:\n",
    "        parents_dct = {i:x for i, x in enumerate(receiver)}\n",
    "        # parents_dct[-1] = None\n",
    "        self.tree_doubling = TreeDoubling(parents=parents_dct, k=10**10)\n",
    "        # node: 要查询的节点, k: 与节点x的距离, kth_node: 距离为k的祖先节点, val: 倍增累加和\n",
    "        ans = 0\n",
    "        for i in range(len(receiver)):\n",
    "            kth_node, val = self.tree_doubling.get_Kth_ancestor(node=i, k=k)\n",
    "            ans = max(ans, val)\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def getMaxFunctionValue(self, receiver: List[int], k: int) -> int:\n",
    "        n = len(receiver)\n",
    "        mx = k.bit_length()\n",
    "        F = [[None] * n for i in range(mx)]\n",
    "        F[0] = [[i,i] for i in receiver]\n",
    "        for i in range(1,mx):\n",
    "            for strt in range(n):\n",
    "                end,v = F[i-1][strt]\n",
    "                F[i][strt] = [F[i-1][end][0],v+F[i-1][end][1]]\n",
    "        res = -inf\n",
    "        for i in range(n):\n",
    "            temp = i;x = k;pos = 0;strt = i\n",
    "            while x:\n",
    "                if x & 1:\n",
    "                    strt,add = F[pos][strt]\n",
    "                    temp += add\n",
    "                pos += 1\n",
    "                x >>= 1\n",
    "            res = max(res,temp)\n",
    "        return res"
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 2
}
